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

llvm#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 llvm#132245,
and saw that the resulting vector and scalar forms were different.

---------

Co-authored-by: Yingwei Zheng <[email protected]>

Author: lukel97

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*.

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)

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
+}

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
+}

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)