[llvm] [InstCombine] Allow folds of shifts by constants for scalable vectors again (PR #132522)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-transforms

Author: Luke Lau (lukel97)

<details>
<summary>Changes</summary>

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

This fixes it by explicitly allowing splats of ConstantInts, it's not the prettiest so open to any suggestions.

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/132522.diff


3 Files Affected:

- (modified) llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp (+4-1) 
- (modified) llvm/test/Transforms/InstCombine/shl-bo.ll (+11) 
- (modified) llvm/test/Transforms/InstCombine/shl-twice-constant.ll (+11) 


``````````diff
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
index 90cd279e8a457..91174cc79cd2b 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp
@@ -428,7 +428,10 @@ Instruction *InstCombinerImpl::commonShiftTransforms(BinaryOperator &I) {
         return R;
 
   Constant *CUI;
-  if (match(Op1, m_ImmConstant(CUI)))
+  if (match(Op1, m_Constant(CUI)) &&
+      (!isa<ConstantExpr>(CUI) ||
+       (Ty->isVectorTy() &&
+        isa_and_present<ConstantInt>(CUI->getSplatValue()))))
     if (Instruction *Res = FoldShiftByConstant(Op0, CUI, I))
       return Res;
 
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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/132522