From 855e02e799b9b05babc39cba1cb121ad7ef50f57 Mon Sep 17 00:00:00 2001
From: Christopher Tetreault <ctetreau@quicinc.com>
Date: Tue, 5 May 2020 14:21:59 -0700
Subject: [PATCH] [SVE] Fix invalid usage of getNumElements() in
 InstCombineMulDivRem

Summary:
getLogBase2 tries to iterate over the number of vector elements. Since
the number of elements of a scalable vector is unknown at compile time,
we must return null if the input type is scalable.

Identified by test LLVM.Transforms/InstCombine::nsw.ll

Reviewers: efriedma, fpetrogalli, kmclaughlin, spatel

Reviewed By: efriedma, fpetrogalli

Subscribers: tschuett, hiraditya, rkruppe, psnobl, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D79197
---
 .../InstCombine/InstCombineMulDivRem.cpp      | 10 ++++---
 .../InstCombine/udiv-pow2-vscale.ll           | 27 +++++++++++++++++++
 2 files changed, 33 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/InstCombine/udiv-pow2-vscale.ll

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
index 5eb281f3c863..775c0e94ab2d 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp
@@ -96,19 +96,21 @@ static Value *simplifyValueKnownNonZero(Value *V, InstCombiner &IC,
 
 /// A helper routine of InstCombiner::visitMul().
 ///
-/// If C is a scalar/vector of known powers of 2, then this function returns
-/// a new scalar/vector obtained from logBase2 of C.
+/// If C is a scalar/fixed width vector of known powers of 2, then this
+/// function returns a new scalar/fixed width vector obtained from logBase2
+/// of C.
 /// Return a null pointer otherwise.
 static Constant *getLogBase2(Type *Ty, Constant *C) {
   const APInt *IVal;
   if (match(C, m_APInt(IVal)) && IVal->isPowerOf2())
     return ConstantInt::get(Ty, IVal->logBase2());
 
-  if (!Ty->isVectorTy())
+  // FIXME: We can extract pow of 2 of splat constant for scalable vectors.
+  if (!isa<FixedVectorType>(Ty))
     return nullptr;
 
   SmallVector<Constant *, 4> Elts;
-  for (unsigned I = 0, E = cast<VectorType>(Ty)->getNumElements(); I != E;
+  for (unsigned I = 0, E = cast<FixedVectorType>(Ty)->getNumElements(); I != E;
        ++I) {
     Constant *Elt = C->getAggregateElement(I);
     if (!Elt)
diff --git a/llvm/test/Transforms/InstCombine/udiv-pow2-vscale.ll b/llvm/test/Transforms/InstCombine/udiv-pow2-vscale.ll
new file mode 100644
index 000000000000..f8f93150417c
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/udiv-pow2-vscale.ll
@@ -0,0 +1,27 @@
+; RUN: opt -instcombine -S < %s | FileCheck %s
+
+; This vscale udiv with a power-of-2 spalt on the rhs should not crash opt
+
+; CHECK: define <vscale x 2 x i32> @udiv_pow2_vscale(<vscale x 2 x i32> %lhs)
+define <vscale x 2 x i32> @udiv_pow2_vscale(<vscale x 2 x i32> %lhs) {
+  %splatter = insertelement <vscale x 2 x i32> undef, i32 2, i32 0
+  %rhs = shufflevector <vscale x 2 x i32> %splatter,
+                       <vscale x 2 x i32> undef,
+                       <vscale x 2 x i32> zeroinitializer
+  %res = udiv <vscale x 2 x i32> %lhs, %rhs
+  ret <vscale x 2 x i32> %res
+}
+
+; This fixed width udiv with a power-of-2 splat on the rhs should also not
+; crash, and instcombine should eliminate the udiv
+
+; CHECK-LABEL: define <2 x i32> @udiv_pow2_fixed(<2 x i32> %lhs)
+; CHECK-NOT: udiv
+define <2 x i32> @udiv_pow2_fixed(<2 x i32> %lhs) {
+  %splatter = insertelement <2 x i32> undef, i32 2, i32 0
+  %rhs = shufflevector <2 x i32> %splatter,
+                       <2 x i32> undef,
+                       <2 x i32> zeroinitializer
+  %res = udiv <2 x i32> %lhs, %rhs
+  ret <2 x i32> %res
+}