[llvm] 3212a08 - [Constant] Allow ConstantAggregateZero a scalable element count

[Fraser Cormack via llvm-commits]

Author: Fraser Cormack
Date: 2021-05-10T13:51:53+01:00
New Revision: 3212a08a8c811441ca68009118758998750ce905

URL: https://github.com/llvm/llvm-project/commit/3212a08a8c811441ca68009118758998750ce905
DIFF: https://github.com/llvm/llvm-project/commit/3212a08a8c811441ca68009118758998750ce905.diff

LOG: [Constant] Allow ConstantAggregateZero a scalable element count

A ConstantAggregateZero may be created from a scalable vector type.
However, it still assumed fixed number of elements when queried for
them. This patch changes ConstantAggregateZero to correctly report its
element count.

This change fixes a couple of issues. Firstly, it fixes a crash in
Constant::getUniqueValue when called on a scalable-vector
zeroinitializer constant.

Secondly, it fixes a latent bug in GlobalISel's IRTranslator in which
translating a scalable-vector zeroinitializer would hit the assertion in
ConstantAggregateZero::getNumElements when casting to a FixedVectorType,
rather than reporting an error more gracefully. This is currently
hypothetical as the IRTranslator has deeper issues preventing the use of
scalable vector types.

Reviewed By: RKSimon

Differential Revision: https://reviews.llvm.org/D102082

Added: 
    llvm/test/Transforms/InstCombine/scalable-select.ll

Modified: 
    llvm/include/llvm/IR/Constants.h
    llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
    llvm/lib/IR/Constants.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/IR/Constants.h b/llvm/include/llvm/IR/Constants.h
index 223e47aa84e7..07d8e9ab5bb6 100644
--- a/llvm/include/llvm/IR/Constants.h
+++ b/llvm/include/llvm/IR/Constants.h
@@ -360,7 +360,7 @@ class ConstantAggregateZero final : public ConstantData {
   Constant *getElementValue(unsigned Idx) const;
 
   /// Return the number of elements in the array, vector, or struct.
-  unsigned getNumElements() const;
+  ElementCount getElementCount() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast:
   ///

diff  --git a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
index 01295a03bae1..ebec73692de8 100644
--- a/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/IRTranslator.cpp
@@ -2876,14 +2876,15 @@ bool IRTranslator::translate(const Constant &C, Register Reg) {
   else if (auto GV = dyn_cast<GlobalValue>(&C))
     EntryBuilder->buildGlobalValue(Reg, GV);
   else if (auto CAZ = dyn_cast<ConstantAggregateZero>(&C)) {
-    if (!CAZ->getType()->isVectorTy())
+    if (!isa<FixedVectorType>(CAZ->getType()))
       return false;
     // Return the scalar if it is a <1 x Ty> vector.
-    if (CAZ->getNumElements() == 1)
+    unsigned NumElts = CAZ->getElementCount().getFixedValue();
+    if (NumElts == 1)
       return translateCopy(C, *CAZ->getElementValue(0u), *EntryBuilder.get());
     SmallVector<Register, 4> Ops;
-    for (unsigned i = 0; i < CAZ->getNumElements(); ++i) {
-      Constant &Elt = *CAZ->getElementValue(i);
+    for (unsigned I = 0; I < NumElts; ++I) {
+      Constant &Elt = *CAZ->getElementValue(I);
       Ops.push_back(getOrCreateVReg(Elt));
     }
     EntryBuilder->buildBuildVector(Reg, Ops);

diff  --git a/llvm/lib/IR/Constants.cpp b/llvm/lib/IR/Constants.cpp
index c0b8af8b2437..7d93d41a398c 100644
--- a/llvm/lib/IR/Constants.cpp
+++ b/llvm/lib/IR/Constants.cpp
@@ -425,13 +425,15 @@ Constant *Constant::getAggregateElement(unsigned Elt) const {
   if (const auto *CC = dyn_cast<ConstantAggregate>(this))
     return Elt < CC->getNumOperands() ? CC->getOperand(Elt) : nullptr;
 
+  if (const auto *CAZ = dyn_cast<ConstantAggregateZero>(this))
+    return Elt < CAZ->getElementCount().getKnownMinValue()
+               ? CAZ->getElementValue(Elt)
+               : nullptr;
+
   // FIXME: getNumElements() will fail for non-fixed vector types.
   if (isa<ScalableVectorType>(getType()))
     return nullptr;
 
-  if (const auto *CAZ = dyn_cast<ConstantAggregateZero>(this))
-    return Elt < CAZ->getNumElements() ? CAZ->getElementValue(Elt) : nullptr;
-
   if (const auto *PV = dyn_cast<PoisonValue>(this))
     return Elt < PV->getNumElements() ? PV->getElementValue(Elt) : nullptr;
 
@@ -1088,13 +1090,13 @@ Constant *ConstantAggregateZero::getElementValue(unsigned Idx) const {
   return getStructElement(Idx);
 }
 
-unsigned ConstantAggregateZero::getNumElements() const {
+ElementCount ConstantAggregateZero::getElementCount() const {
   Type *Ty = getType();
   if (auto *AT = dyn_cast<ArrayType>(Ty))
-    return AT->getNumElements();
+    return ElementCount::getFixed(AT->getNumElements());
   if (auto *VT = dyn_cast<VectorType>(Ty))
-    return cast<FixedVectorType>(VT)->getNumElements();
-  return Ty->getStructNumElements();
+    return VT->getElementCount();
+  return ElementCount::getFixed(Ty->getStructNumElements());
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/llvm/test/Transforms/InstCombine/scalable-select.ll b/llvm/test/Transforms/InstCombine/scalable-select.ll
new file mode 100644
index 000000000000..18370ad76d30
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/scalable-select.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -instcombine -S < %s | FileCheck %s
+
+; This test checks a regression in the select operand folding combine, in which
+; Constant::getUniqueInteger would crash for a scalable-vector zeroinitializer.
+define <vscale x 1 x i32> @select_opt(<vscale x 1 x i32> %b, <vscale x 1 x i1> %m) {
+; CHECK-LABEL: @select_opt(
+; CHECK-NEXT:    [[C:%.*]] = add nsw <vscale x 1 x i32> [[B:%.*]], shufflevector (<vscale x 1 x i32> insertelement (<vscale x 1 x i32> undef, i32 2, i32 0), <vscale x 1 x i32> undef, <vscale x 1 x i32> zeroinitializer)
+; CHECK-NEXT:    [[D:%.*]] = select <vscale x 1 x i1> [[M:%.*]], <vscale x 1 x i32> [[C]], <vscale x 1 x i32> [[B]]
+; CHECK-NEXT:    ret <vscale x 1 x i32> [[D]]
+;
+  %head = insertelement <vscale x 1 x i32> undef, i32 2, i32 0
+  %splat = shufflevector <vscale x 1 x i32> %head, <vscale x 1 x i32> undef, <vscale x 1 x i32> zeroinitializer
+  %c = add nsw <vscale x 1 x i32> %b, %splat
+  %d = select <vscale x 1 x i1> %m, <vscale x 1 x i32> %c, <vscale x 1 x i32> %b
+  ret <vscale x 1 x i32> %d
+}