[PATCH] D103180: [InstSimplify] Add constant fold for extractelement + splat for scalable vectors

[Caroline via Phabricator via llvm-commits]

CarolineConcatto created this revision.
Herald added subscribers: dexonsmith, hiraditya.
CarolineConcatto requested review of this revision.
Herald added a project: LLVM.
Herald added a subscriber: llvm-commits.

This patch allows that scalable vector can fold extractelement and constant splat
only when the lane index is lower than the minimum number of elements of the vector.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D103180

Files:
  llvm/lib/IR/ConstantFold.cpp
  llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll


Index: llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll
===================================================================
--- llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll
+++ llvm/test/Transforms/InstSimplify/ConstProp/extractelement-vscale.ll
@@ -1,4 +1,4 @@
-; RUN: opt -instcombine -S < %s | FileCheck %s
+; RUN: opt -S -instsimplify < %s | FileCheck %s
 
 ; CHECK-LABEL: definitely_in_bounds
 ; CHECK: ret i8 0
@@ -11,3 +11,35 @@
 define i8 @maybe_in_bounds() {
   ret i8 extractelement (<vscale x 16 x i8> zeroinitializer, i64 16)
 }
+
+; Examples of extract a lane from a splat constant
+
+define i32 @extractconstant_shuffle_in_range(i32 %v) {
+; CHECK-LABEL: @extractconstant_shuffle_in_range(
+; CHECK-NEXT:    ret i32 1024
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 1024, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 1
+  ret i32 %r
+}
+
+define i32 @extractconstant_shuffle_maybe_out_of_range(i32 %v) {
+; CHECK-LABEL: @extractconstant_shuffle_maybe_out_of_range(
+; CHECK-NEXT:    ret i32 extractelement (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1024, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer), i32 4)
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 1024, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 4
+  ret i32 %r
+}
+
+define i32 @extractconstant_shuffle_invalid_index(i32 %v) {
+; CHECK-LABEL: @extractconstant_shuffle_invalid_index(
+; CHECK-NEXT:    ret i32 extractelement (<vscale x 4 x i32> shufflevector (<vscale x 4 x i32> insertelement (<vscale x 4 x i32> undef, i32 1024, i32 0), <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer), i32 -1)
+;
+  %in = insertelement <vscale x 4 x i32> undef, i32 1024, i32 0
+  %splat = shufflevector <vscale x 4 x i32> %in, <vscale x 4 x i32> undef, <vscale x 4 x i32> zeroinitializer
+  %r = extractelement <vscale x 4 x i32> %splat, i32 -1
+  ret i32 %r
+}
Index: llvm/lib/IR/ConstantFold.cpp
===================================================================
--- llvm/lib/IR/ConstantFold.cpp
+++ llvm/lib/IR/ConstantFold.cpp
@@ -913,6 +913,8 @@
     if (!CIdx->uge(ValSVTy->getMinNumElements())) {
       if (auto *CAZ = dyn_cast<ConstantAggregateZero>(Val))
         return CAZ->getElementValue(CIdx->getZExtValue());
+      if (Constant *SplatVal = Val->getSplatValue())
+        return SplatVal;
     }
     return nullptr;
   }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D103180.347994.patch
Type: text/x-patch
Size: 2710 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20210526/82ef8dea/attachment.bin>