[PATCH] D103153: [InstCombine] Add fold for extracting known elements from a stepvector

[Caroline via Phabricator via llvm-commits]

CarolineConcatto updated this revision to Diff 347924.
CarolineConcatto marked an inline comment as done.
CarolineConcatto added a comment.

-address Jay Foad's comments


Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D103153/new/

https://reviews.llvm.org/D103153

Files:
  llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
  llvm/test/Transforms/InstCombine/vscale_extractelement.ll


Index: llvm/test/Transforms/InstCombine/vscale_extractelement.ll
===================================================================
--- llvm/test/Transforms/InstCombine/vscale_extractelement.ll
+++ llvm/test/Transforms/InstCombine/vscale_extractelement.ll
@@ -183,3 +183,57 @@
   %E = extractelement <vscale x 4 x i32> %vec_int, i32 2147483647
   ret i32 %E
 }
+
+; Step vector optimization
+
+define i64 @ext_lane0_from_stepvec() {
+; CHECK-LABEL: @ext_lane0_from_stepvec(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret i64 0
+;
+entry:
+  %0 = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %1 = extractelement <vscale x 4 x i64> %0, i32 0
+  ret i64 %1
+}
+
+define i64 @ext_lane_out_of_range_from_stepvec() {
+; CHECK-LABEL: @ext_lane_out_of_range_from_stepvec(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 4 x i64> [[TMP0]], i32 5
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+entry:
+  %0 = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %1 = extractelement <vscale x 4 x i64> %0, i32 5
+  ret i64 %1
+}
+
+define i64 @ext_lane_invalid_from_stepvec() {
+; CHECK-LABEL: @ext_lane_invalid_from_stepvec(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 4 x i64> [[TMP0]], i32 -1
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+entry:
+  %0 = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %1 = extractelement <vscale x 4 x i64> %0, i32 -1
+  ret i64 %1
+}
+
+define i64 @ext_lane_unknown_from_stepvec(i32 %v) {
+; CHECK-LABEL: @ext_lane_unknown_from_stepvec(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+; CHECK-NEXT:    [[TMP1:%.*]] = extractelement <vscale x 4 x i64> [[TMP0]], i32 [[V:%.*]]
+; CHECK-NEXT:    ret i64 [[TMP1]]
+;
+entry:
+  %0 = call <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64()
+  %1 = extractelement <vscale x 4 x i64> %0, i32 %v
+  ret i64 %1
+}
+
+declare <vscale x 4 x i64> @llvm.experimental.stepvector.nxv4i64();
Index: llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
===================================================================
--- llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
+++ llvm/lib/Transforms/InstCombine/InstCombineVectorOps.cpp
@@ -347,6 +347,21 @@
     ElementCount EC = EI.getVectorOperandType()->getElementCount();
     unsigned NumElts = EC.getKnownMinValue();
 
+    if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(SrcVec)) {
+      Intrinsic::ID IID = II->getIntrinsicID();
+      switch (IID) {
+      case Intrinsic::experimental_stepvector:
+        if (IndexC->getValue().ult(NumElts)) {
+          auto *Idx =
+              ConstantInt::get(EI.getVectorOperandType()->getScalarType(),
+                               IndexC->getZExtValue());
+          return replaceInstUsesWith(EI, Idx);
+        }
+        break;
+      default:
+        break;
+      }
+    }
     // InstSimplify should handle cases where the index is invalid.
     // For fixed-length vector, it's invalid to extract out-of-range element.
     if (!EC.isScalable() && IndexC->getValue().uge(NumElts))


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D103153.347924.patch
Type: text/x-patch
Size: 3350 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210526/0767b081/attachment.bin>