Reapply "[InstCombine] Match scalable splats in m_ImmConstant (#132522)"

[lukel97]

This reapplies #132522.

Previously casts of scalable m_ImmConstant splats weren't being folded by ConstantFoldCastOperand, triggering the "Constant-fold of ImmConstant should not fail" assertion.

There are no changes to the code in this PR, instead we just needed #133207 to land first.

A test has been added for the assertion in llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll @icmp_ult_sext_scalable_splat_is_true.

---

#118806 fixed an infinite loop in FoldShiftByConstant that could occur when the shift amount was a ConstantExpr.

However this meant that FoldShiftByConstant no longer kicked in for scalable vectors because scalable splats are represented by ConstantExprs.

This fixes it by allowing scalable splats of non-ConstantExprs in m_ImmConstant, which also fixes a few other test cases where scalable splats were being missed.

But I'm also hoping that UseConstantIntForScalableSplat will eventually remove the need for this.

I noticed this when trying to reverse a combine on RISC-V in #132245, and saw that the resulting vector and scalar forms were different.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-ir

Author: Luke Lau (lukel97)

Changes

This reapplies #132522.

Previously casts of scalable m_ImmConstant splats weren't being folded by ConstantFoldCastOperand, triggering the "Constant-fold of ImmConstant should not fail" assertion.

There are no changes to the code in this PR, instead we just needed #133207 to land first.

A test has been added for the assertion in llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll @icmp_ult_sext_scalable_splat_is_true.

<hr/>

#118806 fixed an infinite loop in FoldShiftByConstant that could occur when the shift amount was a ConstantExpr.

However this meant that FoldShiftByConstant no longer kicked in for scalable vectors because scalable splats are represented by ConstantExprs.

This fixes it by allowing scalable splats of non-ConstantExprs in m_ImmConstant, which also fixes a few other test cases where scalable splats were being missed.

But I'm also hoping that UseConstantIntForScalableSplat will eventually remove the need for this.

I noticed this when trying to reverse a combine on RISC-V in #132245, and saw that the resulting vector and scalar forms were different.

---

Full diff: https://github.com/llvm/llvm-project/pull/134262.diff

6 Files Affected:

• (modified) llvm/include/llvm/IR/PatternMatch.h (+42-9)
• (modified) llvm/test/Transforms/InstCombine/select.ll (+1-2)
• (modified) llvm/test/Transforms/InstCombine/shl-bo.ll (+11)
• (modified) llvm/test/Transforms/InstCombine/shl-twice-constant.ll (+11)
• (modified) llvm/test/Transforms/InstCombine/sub.ll (+1-3)
• (modified) llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll (+9)

diff --git a/llvm/include/llvm/IR/PatternMatch.h b/llvm/include/llvm/IR/PatternMatch.h
index b3eeb1d7ba88a..2d27c19e1b85e 100644
--- a/llvm/include/llvm/IR/PatternMatch.h
+++ b/llvm/include/llvm/IR/PatternMatch.h
@@ -858,18 +858,51 @@ inline bind_ty<const BasicBlock> m_BasicBlock(const BasicBlock *&V) {
   return V;
 }
 
+// TODO: Remove once UseConstant{Int,FP}ForScalableSplat is enabled by default,
+// and use m_Unless(m_ConstantExpr).
+struct immconstant_ty {
+  template <typename ITy> static bool isImmConstant(ITy *V) {
+    if (auto *CV = dyn_cast<Constant>(V)) {
+      if (!isa<ConstantExpr>(CV) && !CV->containsConstantExpression())
+        return true;
+
+      if (CV->getType()->isVectorTy()) {
+        if (auto *Splat = CV->getSplatValue(/*AllowPoison=*/true)) {
+          if (!isa<ConstantExpr>(Splat) &&
+              !Splat->containsConstantExpression()) {
+            return true;
+          }
+        }
+      }
+    }
+    return false;
+  }
+};
+
+struct match_immconstant_ty : immconstant_ty {
+  template <typename ITy> bool match(ITy *V) { return isImmConstant(V); }
+};
+
 /// Match an arbitrary immediate Constant and ignore it.
-inline match_combine_and<class_match<Constant>,
-                         match_unless<constantexpr_match>>
-m_ImmConstant() {
-  return m_CombineAnd(m_Constant(), m_Unless(m_ConstantExpr()));
-}
+inline match_immconstant_ty m_ImmConstant() { return match_immconstant_ty(); }
+
+struct bind_immconstant_ty : immconstant_ty {
+  Constant *&VR;
+
+  bind_immconstant_ty(Constant *&V) : VR(V) {}
+
+  template <typename ITy> bool match(ITy *V) {
+    if (isImmConstant(V)) {
+      VR = cast<Constant>(V);
+      return true;
+    }
+    return false;
+  }
+};
 
 /// Match an immediate Constant, capturing the value if we match.
-inline match_combine_and<bind_ty<Constant>,
-                         match_unless<constantexpr_match>>
-m_ImmConstant(Constant *&C) {
-  return m_CombineAnd(m_Constant(C), m_Unless(m_ConstantExpr()));
+inline bind_immconstant_ty m_ImmConstant(Constant *&C) {
+  return bind_immconstant_ty(C);
 }
 
 /// Match a specified Value*.
diff --git a/llvm/test/Transforms/InstCombine/select.ll b/llvm/test/Transforms/InstCombine/select.ll
index 2078b795817f8..3d81b72dd232e 100644
--- a/llvm/test/Transforms/InstCombine/select.ll
+++ b/llvm/test/Transforms/InstCombine/select.ll
@@ -3519,8 +3519,7 @@ define <vscale x 2 x i32> @scalable_sign_bits(<vscale x 2 x i8> %x) {
 
 define <vscale x 2 x i1> @scalable_non_zero(<vscale x 2 x i32> %x) {
 ; CHECK-LABEL: @scalable_non_zero(
-; CHECK-NEXT:    [[A:%.*]] = or <vscale x 2 x i32> [[X:%.*]], splat (i32 1)
-; CHECK-NEXT:    [[CMP:%.*]] = icmp ult <vscale x 2 x i32> [[A]], splat (i32 57)
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ult <vscale x 2 x i32> [[X:%.*]], splat (i32 56)
 ; CHECK-NEXT:    ret <vscale x 2 x i1> [[CMP]]
 ;
   %a = or <vscale x 2 x i32> %x, splat (i32 1)
diff --git a/llvm/test/Transforms/InstCombine/shl-bo.ll b/llvm/test/Transforms/InstCombine/shl-bo.ll
index c32ac2eacb25a..5ee8716d5d119 100644
--- a/llvm/test/Transforms/InstCombine/shl-bo.ll
+++ b/llvm/test/Transforms/InstCombine/shl-bo.ll
@@ -656,3 +656,14 @@ define <16 x i8> @test_FoldShiftByConstant_CreateAnd(<16 x i8> %in0) {
   %vshl_n = shl <16 x i8> %tmp, <i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5, i8 5>
   ret <16 x i8> %vshl_n
 }
+
+define <vscale x 1 x i8> @test_FoldShiftByConstant_CreateAnd_scalable(<vscale x 1 x i8> %x) {
+; CHECK-LABEL: @test_FoldShiftByConstant_CreateAnd_scalable(
+; CHECK-NEXT:    [[TMP1:%.*]] = shl <vscale x 1 x i8> [[X:%.*]], splat (i8 2)
+; CHECK-NEXT:    [[TMP2:%.*]] = and <vscale x 1 x i8> [[TMP1]], splat (i8 8)
+; CHECK-NEXT:    ret <vscale x 1 x i8> [[TMP2]]
+;
+  %1 = and <vscale x 1 x i8> %x, splat (i8 2)
+  %2 = shl <vscale x 1 x i8> %1, splat (i8 2)
+  ret <vscale x 1 x i8> %2
+}
diff --git a/llvm/test/Transforms/InstCombine/shl-twice-constant.ll b/llvm/test/Transforms/InstCombine/shl-twice-constant.ll
index bbdd7fa3d1c40..151db29fe3e5f 100644
--- a/llvm/test/Transforms/InstCombine/shl-twice-constant.ll
+++ b/llvm/test/Transforms/InstCombine/shl-twice-constant.ll
@@ -14,3 +14,14 @@ define i64 @testfunc() {
   %shl2 = shl i64 %shl1, ptrtoint (ptr @c to i64)
   ret i64 %shl2
 }
+
+define <vscale x 1 x i64> @scalable() {
+; CHECK-LABEL: @scalable(
+; CHECK-NEXT:    [[SHL1:%.*]] = shl nuw <vscale x 1 x i64> splat (i64 1), shufflevector (<vscale x 1 x i64> insertelement (<vscale x 1 x i64> poison, i64 ptrtoint (ptr @c2 to i64), i64 0), <vscale x 1 x i64> poison, <vscale x 1 x i32> zeroinitializer)
+; CHECK-NEXT:    [[SHL2:%.*]] = shl <vscale x 1 x i64> [[SHL1]], shufflevector (<vscale x 1 x i64> insertelement (<vscale x 1 x i64> poison, i64 ptrtoint (ptr @c to i64), i64 0), <vscale x 1 x i64> poison, <vscale x 1 x i32> zeroinitializer)
+; CHECK-NEXT:    ret <vscale x 1 x i64> [[SHL2]]
+;
+  %shl1 = shl <vscale x 1 x i64> splat (i64 1), splat (i64 ptrtoint (ptr @c2 to i64))
+  %shl2 = shl <vscale x 1 x i64> %shl1, splat (i64 ptrtoint (ptr @c to i64))
+  ret <vscale x 1 x i64> %shl2
+}
diff --git a/llvm/test/Transforms/InstCombine/sub.ll b/llvm/test/Transforms/InstCombine/sub.ll
index e89419d1f3838..81ecd8506514e 100644
--- a/llvm/test/Transforms/InstCombine/sub.ll
+++ b/llvm/test/Transforms/InstCombine/sub.ll
@@ -857,11 +857,9 @@ define <2 x i16> @test44vecminval(<2 x i16> %x) {
   ret <2 x i16> %sub
 }
 
-; FIXME: This isn't combined to xor as above because the pattern in visitSub
-; uses m_ImmConstant which matches Constant but (explicitly) not ConstantExpr.
 define <vscale x 2 x i16> @test44scalablevecminval(<vscale x 2 x i16> %x) {
 ; CHECK-LABEL: @test44scalablevecminval(
-; CHECK-NEXT:    [[SUB:%.*]] = add <vscale x 2 x i16> [[X:%.*]], splat (i16 -32768)
+; CHECK-NEXT:    [[SUB:%.*]] = xor <vscale x 2 x i16> [[X:%.*]], splat (i16 -32768)
 ; CHECK-NEXT:    ret <vscale x 2 x i16> [[SUB]]
 ;
   %sub = sub nsw <vscale x 2 x i16> %x, splat (i16 -32768)
diff --git a/llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll b/llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll
index 8b27ab1f0ef26..3314ef6ff9715 100644
--- a/llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll
+++ b/llvm/test/Transforms/InstSimplify/vec-icmp-of-cast.ll
@@ -174,3 +174,12 @@ define <2 x i1> @icmp_slt_sext_is_true_false(<2 x i8> %x) {
   %cmp = icmp slt <2 x i32> %xext, <i32 257, i32 -450>
   ret <2 x i1> %cmp
 }
+
+define <vscale x 4 x i1> @icmp_ult_sext_scalable_splat_is_true(<vscale x 4 x i8> %x) {
+; CHECK-LABEL: @icmp_ult_sext_scalable_splat_is_true(
+; CHECK-NEXT:    ret <vscale x 4 x i1> splat (i1 true)
+;
+  %s = sext <vscale x 4 x i8> %x to <vscale x 4 x i64>
+  %cmp = icmp slt <vscale x 4 x i64> %s, splat (i64 257)
+  ret <vscale x 4 x i1> %cmp
+}