[Scalarizer] InsertElement handling w/ variable insert index (PR46524)

[lebedev.ri]

Summary

I'm interested in taking the original C++ input,
for which we currently are stuck with an alloca
and producing roughly the lower IR,
with neither an alloca nor a vector ops:
https://godbolt.org/z/cRRWaJ

For that, as intermediate step, i'd to somehow perform scalarization.
As per @arsenmn suggestion, i'm trying to see if scalarizer can help me
avoid writing a bicycle.

I'm not sure if it's really intentional that variable insert is not handled currently.
If it really is, and is supposed to stay that way (?), i guess i could guard it..

See PR46524.

[bjope]

I should say something, as we use the pass downstream (we run Scalarizer early in llc pipeline, partly for historical reasons, but still due to sometimes getting better code when scalarizing before some other pass (maybe LoopStrengthReduce/CodeGenPrepare?) instead of doing it at ISel).

I think the normal expand at ISel is doing a simple, store vector to stack, do an indexed store, load back the vector. So one could get condition/branch free code (except a check that index is within range). If such a lowering is good/bad for a target will ofcourse depend on if load/store is fast etc (might depend on number of elements in the vector as well). Basically it is easier to deal with the vector version and convert it to whatever is suitable for the target at ISel (rather than trying to convert the scalarized sequence into something using the indexed write technique). No idea if that is the intent with not rewriting insertelement/extractelement with variable index. But this patch might hurt someone out there that is expecting such lowering.

With the above said:

• I've checked our downstream lowering (at ISel) and we do lower the insertelement using the standard expand technique (we get a vector store, an indexed store for the element and a vector load). With this patch we get code with lots of conditional moves. However, in reality these insertelements with variable index probably never happens so it doesn't really matter what we do.
• With D82970 in mind, our downstream lowering (at ISel) is expanding an extractelement by using the conditional moves. So that patch has even less impact for us.
• The pass is supposed to remove vector ops, so I don't think it is essentially wrong to do these new rewrites. However, it still using insertelement/extractelement with constant indices so it is not obvious that some insert/extracts should be scalarized and others not.
• I don't think it hurts to add an option that could toggle these new scalarizations on/off (in case someone runs into trouble with this, considering that the pass hasn't scalarized these in the past . Then at least, someone downstream could easily benchmark what happens when toggling this option. Or even override the default somehow if that is desired.

No expert on poison etc, but I figure the resulting code is less poisonous (insertelement can generate poison, but the select sequence won't). I guess that is ok.

[jdoerfert]

I like the flag idea, but I would expect this to be the default. If I read scalarizer and I enable it, I would expect the code to be scalarized, not partially scalarized based on a heuristic that may or may not be to my liking. The option to opt-out of some scalarization is sensible though.

Out of curiosity, did u try we fold this properly if the index is substituted afterwards with a constant?

I figure the resulting code is less poisonous (insertelement can generate poison, but the select sequence won't). I guess that is ok.

Agreed.

llvm/lib/Transforms/Scalar/Scalarizer.cpp

762

I do not understand why we skip constant indices.

[bjope]

llvm/lib/Transforms/Scalar/Scalarizer.cpp

762

I don't know exactly, but Scalarizer is sometimes creating such instructions itself when a vector still is needed.

(Downstream we have hacked the scalarizer to keep vector load/stores (typically loading/storing to a register pair/quad). But artithmetics etc needs to be scalarized. So there are typically lots of insertelement/extractelement instructions around to glue things together after running the scalarizer.)

[arsenm]

In D82961#2128464, @jdoerfert wrote:

I like the flag idea, but I would expect this to be the default. If I read scalarizer and I enable it, I would expect the code to be scalarized, not partially scalarized based on a heuristic that may or may not be to my liking. The option to opt-out of some scalarization is sensible though.

You have a different idea of what scalarization means than I do. I want to eliminate operations on vectors; I consider this as arithmetic operations operating on vectors. An insertelement/extractelement is not really a vector operation itself.

For AMDGPU we can implement dynamic indexing by directly indexing the registers without splitting this into a long chain of operations or introducing stack use

[lebedev.ri]

Thank you for taking a look everyone!
Adding opt-out flag.

llvm/lib/Transforms/Scalar/Scalarizer.cpp

762

I mean, why shouldn't we skip them.
I'm not sure replacing constant extract with n constant extracts+select is better.

[jdoerfert]

In D82961#2128678, @arsenm wrote:

In D82961#2128464, @jdoerfert wrote:

I like the flag idea, but I would expect this to be the default. If I read scalarizer and I enable it, I would expect the code to be scalarized, not partially scalarized based on a heuristic that may or may not be to my liking. The option to opt-out of some scalarization is sensible though.

You have a different idea of what scalarization means than I do. I want to eliminate operations on vectors; I consider this as arithmetic operations operating on vectors. An insertelement/extractelement is not really a vector operation itself.

We might. My idea of scalarization is that it removes vector operations (= operations interacting with vector values). Put differently, I would expect all types/values to be scalars and not vectors after I run the scalarizer.

llvm/lib/Transforms/Scalar/Scalarizer.cpp

762

Do we need n extracts+select for a constant index? wouldn't that fold into a single value?

[lebedev.ri]

Reordered, precommitted tests.

[jdoerfert]

LGTM

llvm/lib/Transforms/Scalar/Scalarizer.cpp

 #include <iterator>
 #include <map>
 #include <utility>
 
 using namespace llvm;
 
 #define DEBUG_TYPE "scalarizer"
 
+static cl::opt<bool> ScalarizeVariableInsertExtract(
+    "scalarize-variable-insert-extract", cl::init(true), cl::Hidden,
+    cl::desc("Allow the scalarizer pass to scalarize "
+             "insertelement/extractelement with variable index"));
+
 // This is disabled by default because having separate loads and stores
 // makes it more likely that the -combiner-alias-analysis limits will be
 // reached.
 static cl::opt<bool>
     ScalarizeLoadStore("scalarize-load-store", cl::init(false), cl::Hidden,
                        cl::desc("Allow the scalarizer pass to scalarize loads and store"));
 
 namespace {
bool ScalarizerVisitor::visitInsertElementInst(InsertElementInst &IEI) {
 
   unsigned NumElems = VT->getNumElements();
   IRBuilder<> Builder(&IEI);
   Scatterer Op0 = scatter(&IEI, IEI.getOperand(0));
   Value *NewElt = IEI.getOperand(1);
   Value *InsIdx = IEI.getOperand(2);
 
   ValueVector Res;
   Res.resize(NumElems);
 
   if (auto *CI = dyn_cast<ConstantInt>(InsIdx)) {
     for (unsigned I = 0; I < NumElems; ++I)
       Res[I] = CI->getValue().getZExtValue() == I ? NewElt : Op0[I];
   } else {
+    if (!ScalarizeVariableInsertExtract)
       return false;
+
+    for (unsigned I = 0; I < NumElems; ++I) {
+      Res[I] = Builder.CreateSelect(
+          Builder.CreateICmpEQ(InsIdx, ConstantInt::get(InsIdx->getType(), I),
+                               InsIdx->getName() + ".is." + Twine(I)),
+          NewElt, Op0[I], IEI.getName() + ".i" + Twine(I));
+    }
   }
 
   gather(&IEI, Res);
   return true;
 }
 
 bool ScalarizerVisitor::visitExtractElementInst(ExtractElementInst &EEI) {
   VectorType *VT = dyn_cast<VectorType>(EEI.getOperand(0)->getType());

llvm/test/Transforms/Scalarizer/variable-insertelement.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt %s -scalarizer -scalarize-load-store -dce -S | FileCheck --check-prefixes=ALL %s
+; RUN: opt %s -scalarizer -dce -S | FileCheck --check-prefixes=ALL,DEFAULT %s
+; RUN: opt %s -scalarizer -scalarize-variable-insert-extract=false -dce -S | FileCheck --check-prefixes=ALL,OFF %s
+; RUN: opt %s -scalarizer -scalarize-variable-insert-extract=true -dce -S | FileCheck --check-prefixes=ALL,DEFAULT,ON %s
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 
 ; Test that variable inserts are scalarized.
 define <4 x i32> @f1(<4 x i32> %src, i32 %val, i32 %index) {
-; ALL-LABEL: @f1(
-; ALL-NEXT:    [[RES:%.*]] = insertelement <4 x i32> [[SRC:%.*]], i32 [[VAL:%.*]], i32 [[INDEX:%.*]]
-; ALL-NEXT:    ret <4 x i32> [[RES]]
+; DEFAULT-LABEL: @f1(
+; DEFAULT-NEXT:    [[INDEX_IS_0:%.*]] = icmp eq i32 [[INDEX:%.*]], 0
+; DEFAULT-NEXT:    [[SRC_I0:%.*]] = extractelement <4 x i32> [[SRC:%.*]], i32 0
+; DEFAULT-NEXT:    [[RES_I0:%.*]] = select i1 [[INDEX_IS_0]], i32 [[VAL:%.*]], i32 [[SRC_I0]]
+; DEFAULT-NEXT:    [[INDEX_IS_1:%.*]] = icmp eq i32 [[INDEX]], 1
+; DEFAULT-NEXT:    [[SRC_I1:%.*]] = extractelement <4 x i32> [[SRC]], i32 1
+; DEFAULT-NEXT:    [[RES_I1:%.*]] = select i1 [[INDEX_IS_1]], i32 [[VAL]], i32 [[SRC_I1]]
+; DEFAULT-NEXT:    [[INDEX_IS_2:%.*]] = icmp eq i32 [[INDEX]], 2
+; DEFAULT-NEXT:    [[SRC_I2:%.*]] = extractelement <4 x i32> [[SRC]], i32 2
+; DEFAULT-NEXT:    [[RES_I2:%.*]] = select i1 [[INDEX_IS_2]], i32 [[VAL]], i32 [[SRC_I2]]
+; DEFAULT-NEXT:    [[INDEX_IS_3:%.*]] = icmp eq i32 [[INDEX]], 3
+; DEFAULT-NEXT:    [[SRC_I3:%.*]] = extractelement <4 x i32> [[SRC]], i32 3
+; DEFAULT-NEXT:    [[RES_I3:%.*]] = select i1 [[INDEX_IS_3]], i32 [[VAL]], i32 [[SRC_I3]]
+; DEFAULT-NEXT:    [[RES_UPTO0:%.*]] = insertelement <4 x i32> undef, i32 [[RES_I0]], i32 0
+; DEFAULT-NEXT:    [[RES_UPTO1:%.*]] = insertelement <4 x i32> [[RES_UPTO0]], i32 [[RES_I1]], i32 1
+; DEFAULT-NEXT:    [[RES_UPTO2:%.*]] = insertelement <4 x i32> [[RES_UPTO1]], i32 [[RES_I2]], i32 2
+; DEFAULT-NEXT:    [[RES:%.*]] = insertelement <4 x i32> [[RES_UPTO2]], i32 [[RES_I3]], i32 3
+; DEFAULT-NEXT:    ret <4 x i32> [[RES]]
+;
+; OFF-LABEL: @f1(
+; OFF-NEXT:    [[RES:%.*]] = insertelement <4 x i32> [[SRC:%.*]], i32 [[VAL:%.*]], i32 [[INDEX:%.*]]
+; OFF-NEXT:    ret <4 x i32> [[RES]]
 ;
   %res = insertelement <4 x i32> %src, i32 %val, i32 %index
   ret <4 x i32> %res
 }
 
 define void @f2(<4 x i32> *%dest, <4 x i32> *%src, i32 %index) {
-; ALL-LABEL: @f2(
-; ALL-NEXT:    [[DEST_I0:%.*]] = bitcast <4 x i32>* [[DEST:%.*]] to i32*
-; ALL-NEXT:    [[DEST_I1:%.*]] = getelementptr i32, i32* [[DEST_I0]], i32 1
-; ALL-NEXT:    [[DEST_I2:%.*]] = getelementptr i32, i32* [[DEST_I0]], i32 2
-; ALL-NEXT:    [[DEST_I3:%.*]] = getelementptr i32, i32* [[DEST_I0]], i32 3
-; ALL-NEXT:    [[SRC_I0:%.*]] = bitcast <4 x i32>* [[SRC:%.*]] to i32*
-; ALL-NEXT:    [[VAL0_I0:%.*]] = load i32, i32* [[SRC_I0]], align 16
-; ALL-NEXT:    [[SRC_I1:%.*]] = getelementptr i32, i32* [[SRC_I0]], i32 1
-; ALL-NEXT:    [[VAL0_I1:%.*]] = load i32, i32* [[SRC_I1]], align 4
-; ALL-NEXT:    [[SRC_I2:%.*]] = getelementptr i32, i32* [[SRC_I0]], i32 2
-; ALL-NEXT:    [[VAL0_I2:%.*]] = load i32, i32* [[SRC_I2]], align 8
-; ALL-NEXT:    [[SRC_I3:%.*]] = getelementptr i32, i32* [[SRC_I0]], i32 3
-; ALL-NEXT:    [[VAL0_I3:%.*]] = load i32, i32* [[SRC_I3]], align 4
-; ALL-NEXT:    [[VAL0_UPTO0:%.*]] = insertelement <4 x i32> undef, i32 [[VAL0_I0]], i32 0
-; ALL-NEXT:    [[VAL0_UPTO1:%.*]] = insertelement <4 x i32> [[VAL0_UPTO0]], i32 [[VAL0_I1]], i32 1
-; ALL-NEXT:    [[VAL0_UPTO2:%.*]] = insertelement <4 x i32> [[VAL0_UPTO1]], i32 [[VAL0_I2]], i32 2
-; ALL-NEXT:    [[VAL0:%.*]] = insertelement <4 x i32> [[VAL0_UPTO2]], i32 [[VAL0_I3]], i32 3
-; ALL-NEXT:    [[VAL1:%.*]] = insertelement <4 x i32> [[VAL0]], i32 1, i32 [[INDEX:%.*]]
-; ALL-NEXT:    [[VAL1_I0:%.*]] = extractelement <4 x i32> [[VAL1]], i32 0
-; ALL-NEXT:    [[VAL2_I0:%.*]] = shl i32 1, [[VAL1_I0]]
-; ALL-NEXT:    [[VAL1_I1:%.*]] = extractelement <4 x i32> [[VAL1]], i32 1
-; ALL-NEXT:    [[VAL2_I1:%.*]] = shl i32 2, [[VAL1_I1]]
-; ALL-NEXT:    [[VAL1_I2:%.*]] = extractelement <4 x i32> [[VAL1]], i32 2
-; ALL-NEXT:    [[VAL2_I2:%.*]] = shl i32 3, [[VAL1_I2]]
-; ALL-NEXT:    [[VAL1_I3:%.*]] = extractelement <4 x i32> [[VAL1]], i32 3
-; ALL-NEXT:    [[VAL2_I3:%.*]] = shl i32 4, [[VAL1_I3]]
-; ALL-NEXT:    store i32 [[VAL2_I0]], i32* [[DEST_I0]], align 16
-; ALL-NEXT:    store i32 [[VAL2_I1]], i32* [[DEST_I1]], align 4
-; ALL-NEXT:    store i32 [[VAL2_I2]], i32* [[DEST_I2]], align 8
-; ALL-NEXT:    store i32 [[VAL2_I3]], i32* [[DEST_I3]], align 4
-; ALL-NEXT:    ret void
+; DEFAULT-LABEL: @f2(
+; DEFAULT-NEXT:    [[VAL0:%.*]] = load <4 x i32>, <4 x i32>* [[SRC:%.*]], align 16
+; DEFAULT-NEXT:    [[INDEX_IS_0:%.*]] = icmp eq i32 [[INDEX:%.*]], 0
+; DEFAULT-NEXT:    [[VAL0_I0:%.*]] = extractelement <4 x i32> [[VAL0]], i32 0
+; DEFAULT-NEXT:    [[VAL1_I0:%.*]] = select i1 [[INDEX_IS_0]], i32 1, i32 [[VAL0_I0]]
+; DEFAULT-NEXT:    [[INDEX_IS_1:%.*]] = icmp eq i32 [[INDEX]], 1
+; DEFAULT-NEXT:    [[VAL0_I1:%.*]] = extractelement <4 x i32> [[VAL0]], i32 1
+; DEFAULT-NEXT:    [[VAL1_I1:%.*]] = select i1 [[INDEX_IS_1]], i32 1, i32 [[VAL0_I1]]
+; DEFAULT-NEXT:    [[INDEX_IS_2:%.*]] = icmp eq i32 [[INDEX]], 2
+; DEFAULT-NEXT:    [[VAL0_I2:%.*]] = extractelement <4 x i32> [[VAL0]], i32 2
+; DEFAULT-NEXT:    [[VAL1_I2:%.*]] = select i1 [[INDEX_IS_2]], i32 1, i32 [[VAL0_I2]]
+; DEFAULT-NEXT:    [[INDEX_IS_3:%.*]] = icmp eq i32 [[INDEX]], 3
+; DEFAULT-NEXT:    [[VAL0_I3:%.*]] = extractelement <4 x i32> [[VAL0]], i32 3
+; DEFAULT-NEXT:    [[VAL1_I3:%.*]] = select i1 [[INDEX_IS_3]], i32 1, i32 [[VAL0_I3]]
+; DEFAULT-NEXT:    [[VAL2_I0:%.*]] = shl i32 1, [[VAL1_I0]]
+; DEFAULT-NEXT:    [[VAL2_I1:%.*]] = shl i32 2, [[VAL1_I1]]
+; DEFAULT-NEXT:    [[VAL2_I2:%.*]] = shl i32 3, [[VAL1_I2]]
+; DEFAULT-NEXT:    [[VAL2_I3:%.*]] = shl i32 4, [[VAL1_I3]]
+; DEFAULT-NEXT:    [[VAL2_UPTO0:%.*]] = insertelement <4 x i32> undef, i32 [[VAL2_I0]], i32 0
+; DEFAULT-NEXT:    [[VAL2_UPTO1:%.*]] = insertelement <4 x i32> [[VAL2_UPTO0]], i32 [[VAL2_I1]], i32 1
+; DEFAULT-NEXT:    [[VAL2_UPTO2:%.*]] = insertelement <4 x i32> [[VAL2_UPTO1]], i32 [[VAL2_I2]], i32 2
+; DEFAULT-NEXT:    [[VAL2:%.*]] = insertelement <4 x i32> [[VAL2_UPTO2]], i32 [[VAL2_I3]], i32 3
+; DEFAULT-NEXT:    store <4 x i32> [[VAL2]], <4 x i32>* [[DEST:%.*]], align 16
+; DEFAULT-NEXT:    ret void
+;
+; OFF-LABEL: @f2(
+; OFF-NEXT:    [[VAL0:%.*]] = load <4 x i32>, <4 x i32>* [[SRC:%.*]], align 16
+; OFF-NEXT:    [[VAL1:%.*]] = insertelement <4 x i32> [[VAL0]], i32 1, i32 [[INDEX:%.*]]
+; OFF-NEXT:    [[VAL1_I0:%.*]] = extractelement <4 x i32> [[VAL1]], i32 0
+; OFF-NEXT:    [[VAL2_I0:%.*]] = shl i32 1, [[VAL1_I0]]
+; OFF-NEXT:    [[VAL1_I1:%.*]] = extractelement <4 x i32> [[VAL1]], i32 1
+; OFF-NEXT:    [[VAL2_I1:%.*]] = shl i32 2, [[VAL1_I1]]
+; OFF-NEXT:    [[VAL1_I2:%.*]] = extractelement <4 x i32> [[VAL1]], i32 2
+; OFF-NEXT:    [[VAL2_I2:%.*]] = shl i32 3, [[VAL1_I2]]
+; OFF-NEXT:    [[VAL1_I3:%.*]] = extractelement <4 x i32> [[VAL1]], i32 3
+; OFF-NEXT:    [[VAL2_I3:%.*]] = shl i32 4, [[VAL1_I3]]
+; OFF-NEXT:    [[VAL2_UPTO0:%.*]] = insertelement <4 x i32> undef, i32 [[VAL2_I0]], i32 0
+; OFF-NEXT:    [[VAL2_UPTO1:%.*]] = insertelement <4 x i32> [[VAL2_UPTO0]], i32 [[VAL2_I1]], i32 1
+; OFF-NEXT:    [[VAL2_UPTO2:%.*]] = insertelement <4 x i32> [[VAL2_UPTO1]], i32 [[VAL2_I2]], i32 2
+; OFF-NEXT:    [[VAL2:%.*]] = insertelement <4 x i32> [[VAL2_UPTO2]], i32 [[VAL2_I3]], i32 3
+; OFF-NEXT:    store <4 x i32> [[VAL2]], <4 x i32>* [[DEST:%.*]], align 16
+; OFF-NEXT:    ret void
 ;
   %val0 = load <4 x i32> , <4 x i32> *%src
   %val1 = insertelement <4 x i32> %val0, i32 1, i32 %index
   %val2 = shl <4 x i32> <i32 1, i32 2, i32 3, i32 4>, %val1
   store <4 x i32> %val2, <4 x i32> *%dest
   ret void
 }