[Scalarizer] No need to gather a scattered extracted element

[serge-sans-paille]

Summary

ExtractElement does not produce a vector out of a vector, so there's no need to
call a gather once done.

Fix #54469

Credits to [e-mail] for the original approach.

[serge-sans-paille]

Gentle ping :-)

[nikic]

There is another patch for this issue in D123684. Haven't looked at either variant yet, but the approach there looks quite different.

[nikic]

My current thinking here is that we shouldn't be making this gather() call at all. The code seems to be effectively treating the result of the extractelement as a 1-vector, which doesn't make a lot of sense to me.

I tried the following patch, that just replaces the extractelement with the scattered value: https://gist.github.com/nikic/d09ace51902e03e5b0cf04c528f87c49 Apart from one dead instruction, this seems to be working fine with tests.

What do you think?

llvm/test/Transforms/Scalarizer/global-bug.ll

2

Spurious change :)

[serge-sans-paille]

Follow @nikic approach there, it's clean and simple. Just fix a little edge case in the finish routine. Note that I could have checked the size of the cleanup vector, but I felt like a more explicit approach would be better.

[nikic]

LGTM

llvm/test/Transforms/Scalarizer/vector-of-pointer-to-vector.ll

6

Maybe add one more test with constant index, as it uses a different code path in the implementation?

llvm/lib/Transforms/Scalar/Scalarizer.cpp

public:
   bool visitPHINode(PHINode &PHI);
   bool visitLoadInst(LoadInst &LI);
   bool visitStoreInst(StoreInst &SI);
   bool visitCallInst(CallInst &ICI);
 
 private:
   Scatterer scatter(Instruction *Point, Value *V, Type *PtrElemTy = nullptr);
   void gather(Instruction *Op, const ValueVector &CV);
+  void replaceUses(Instruction *Op, Value *CV);
   bool canTransferMetadata(unsigned Kind);
   void transferMetadataAndIRFlags(Instruction *Op, const ValueVector &CV);
   Optional<VectorLayout> getVectorLayout(Type *Ty, Align Alignment,
                                          const DataLayout &DL);
   bool finish();
 
   template<typename T> bool splitUnary(Instruction &, const T &);
   template<typename T> bool splitBinary(Instruction &, const T &);
 
   bool splitCall(CallInst &CI);
 
   ScatterMap Scattered;
   GatherList Gathered;
+  bool Scalarized;
 
   SmallVector<WeakTrackingVH, 32> PotentiallyDeadInstrs;
 
   unsigned ParallelLoopAccessMDKind;
 
   DominatorTree *DT;
 
   const bool ScalarizeVariableInsertExtract;

 
 FunctionPass *llvm::createScalarizerPass() {
   return new ScalarizerLegacyPass();
 }
 
 bool ScalarizerVisitor::visit(Function &F) {
   assert(Gathered.empty() && Scattered.empty());
 
+  Scalarized = false;
+
   // To ensure we replace gathered components correctly we need to do an ordered
   // traversal of the basic blocks in the function.
   ReversePostOrderTraversal<BasicBlock *> RPOT(&F.getEntryBlock());
   for (BasicBlock *BB : RPOT) {
     for (BasicBlock::iterator II = BB->begin(), IE = BB->end(); II != IE;) {
       Instruction *I = &*II;
       bool Done = InstVisitor::visit(I);
       ++II;
for (unsigned I = 0, E = SV.size(); I != E; ++I) {
       Old->replaceAllUsesWith(CV[I]);
       PotentiallyDeadInstrs.emplace_back(Old);
     }
   }
   SV = CV;
   Gathered.push_back(GatherList::value_type(Op, &SV));
 }
 
+// Replace Op with CV and collect Op has a potentially dead instruction.
+void ScalarizerVisitor::replaceUses(Instruction *Op, Value *CV) {
+  if (CV != Op) {
+    Op->replaceAllUsesWith(CV);
+    PotentiallyDeadInstrs.emplace_back(Op);
+    Scalarized = true;
+  }
+}
+
 // Return true if it is safe to transfer the given metadata tag from
 // vector to scalar instructions.
 bool ScalarizerVisitor::canTransferMetadata(unsigned Tag) {
   return (Tag == LLVMContext::MD_tbaa
           || Tag == LLVMContext::MD_fpmath
           || Tag == LLVMContext::MD_tbaa_struct
           || Tag == LLVMContext::MD_invariant_load
           || Tag == LLVMContext::MD_alias_scope
bool ScalarizerVisitor::visitExtractElementInst(ExtractElementInst &EEI) {
 
   unsigned NumSrcElems = cast<FixedVectorType>(VT)->getNumElements();
   IRBuilder<> Builder(&EEI);
   Scatterer Op0 = scatter(&EEI, EEI.getOperand(0));
   Value *ExtIdx = EEI.getOperand(1);
 
   if (auto *CI = dyn_cast<ConstantInt>(ExtIdx)) {
     Value *Res = Op0[CI->getValue().getZExtValue()];
-    gather(&EEI, {Res});
+    replaceUses(&EEI, Res);
     return true;
   }
 
   if (!ScalarizeVariableInsertExtract)
     return false;
 
   Value *Res = UndefValue::get(VT->getElementType());
   for (unsigned I = 0; I < NumSrcElems; ++I) {
     Value *ShouldExtract =
         Builder.CreateICmpEQ(ExtIdx, ConstantInt::get(ExtIdx->getType(), I),
                              ExtIdx->getName() + ".is." + Twine(I));
     Value *Elt = Op0[I];
     Res = Builder.CreateSelect(ShouldExtract, Elt, Res,
                                EEI.getName() + ".upto" + Twine(I));
   }
-  gather(&EEI, {Res});
+  replaceUses(&EEI, Res);
   return true;
 }
 
 bool ScalarizerVisitor::visitShuffleVectorInst(ShuffleVectorInst &SVI) {
   VectorType *VT = dyn_cast<VectorType>(SVI.getType());
   if (!VT)
     return false;
 
bool ScalarizerVisitor::visitCallInst(CallInst &CI) {
   return splitCall(CI);
 }
 
 // Delete the instructions that we scalarized.  If a full vector result
 // is still needed, recreate it using InsertElements.
 bool ScalarizerVisitor::finish() {
   // The presence of data in Gathered or Scattered indicates changes
   // made to the Function.
-  if (Gathered.empty() && Scattered.empty())
+  if (Gathered.empty() && Scattered.empty() && !Scalarized)
     return false;
   for (const auto &GMI : Gathered) {
     Instruction *Op = GMI.first;
     ValueVector &CV = *GMI.second;
     if (!Op->use_empty()) {
       // The value is still needed, so recreate it using a series of
       // InsertElements.
       Value *Res = PoisonValue::get(Op->getType());
if (!Op->use_empty()) {
           continue;
       }
       Op->replaceAllUsesWith(Res);
     }
     PotentiallyDeadInstrs.emplace_back(Op);
   }
   Gathered.clear();
   Scattered.clear();
+  Scalarized = false;
 
   RecursivelyDeleteTriviallyDeadInstructionsPermissive(PotentiallyDeadInstrs);
 
   return true;
 }
 
 PreservedAnalyses ScalarizerPass::run(Function &F, FunctionAnalysisManager &AM) {
   Module &M = *F.getParent();

llvm/test/Transforms/Scalarizer/global-bug.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; RUN: opt < %s -passes='function(scalarizer)' -S | FileCheck %s
 
 @a = dso_local global i16 0, align 1
 @b = dso_local local_unnamed_addr global i16 0, align 1
 
 ; The scalarizer used to take the name of the extractelement instruction
 ; ("extract") and put that on the extracted value, which in this test case is
 ; the global variable @a. That was wrong, as we must not change the name of
 ; the global variable. So make sure we find "@a" in the ptrtoint.

llvm/test/Transforms/Scalarizer/vector-of-pointer-to-vector.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt %s -passes='function(scalarizer,dce)' -scalarize-load-store -S | FileCheck %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"
+
+define <1 x i32> @f1(<1 x <1 x i32>*> %src, i32 %index) {
+; CHECK-LABEL: @f1(
+; CHECK-NEXT:    [[INDEX_IS_0:%.*]] = icmp eq i32 [[INDEX:%.*]], 0
+; CHECK-NEXT:    [[SRC_I0:%.*]] = extractelement <1 x <1 x i32>*> [[SRC:%.*]], i32 0
+; CHECK-NEXT:    [[DOTUPTO0:%.*]] = select i1 [[INDEX_IS_0]], <1 x i32>* [[SRC_I0]], <1 x i32>* undef
+; CHECK-NEXT:    [[DOTUPTO0_I0:%.*]] = bitcast <1 x i32>* [[DOTUPTO0]] to i32*
+; CHECK-NEXT:    [[DOTI0:%.*]] = load i32, i32* [[DOTUPTO0_I0]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <1 x i32> poison, i32 [[DOTI0]], i32 0
+; CHECK-NEXT:    ret <1 x i32> [[TMP1]]
+;
+  %1 = extractelement <1 x <1 x i32>*> %src, i32 %index
+  %2 = load <1 x i32>, <1 x i32>* %1, align 4
+  ret <1 x i32> %2
+}
+
+define <1 x i32> @f1b(<1 x <1 x i32>*> %src) {
+; CHECK-LABEL: @f1b(
+; CHECK-NEXT:    [[SRC_I0:%.*]] = extractelement <1 x <1 x i32>*> [[SRC:%.*]], i32 0
+; CHECK-NEXT:    [[SRC_I0_I0:%.*]] = bitcast <1 x i32>* [[SRC_I0]] to i32*
+; CHECK-NEXT:    [[DOTI0:%.*]] = load i32, i32* [[SRC_I0_I0]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <1 x i32> poison, i32 [[DOTI0]], i32 0
+; CHECK-NEXT:    ret <1 x i32> [[TMP1]]
+;
+  %1 = extractelement <1 x <1 x i32>*> %src, i32 0
+  %2 = load <1 x i32>, <1 x i32>* %1, align 4
+  ret <1 x i32> %2
+}
+
+define <2 x i32> @f2(<1 x <2 x i32>*> %src, i32 %index) {
+; CHECK-LABEL: @f2(
+; CHECK-NEXT:    [[INDEX_IS_0:%.*]] = icmp eq i32 [[INDEX:%.*]], 0
+; CHECK-NEXT:    [[SRC_I0:%.*]] = extractelement <1 x <2 x i32>*> [[SRC:%.*]], i32 0
+; CHECK-NEXT:    [[DOTUPTO0:%.*]] = select i1 [[INDEX_IS_0]], <2 x i32>* [[SRC_I0]], <2 x i32>* undef
+; CHECK-NEXT:    [[DOTUPTO0_I0:%.*]] = bitcast <2 x i32>* [[DOTUPTO0]] to i32*
+; CHECK-NEXT:    [[DOTUPTO0_I1:%.*]] = getelementptr i32, i32* [[DOTUPTO0_I0]], i32 1
+; CHECK-NEXT:    [[DOTI0:%.*]] = load i32, i32* [[DOTUPTO0_I0]], align 4
+; CHECK-NEXT:    [[DOTI1:%.*]] = load i32, i32* [[DOTUPTO0_I1]], align 4
+; CHECK-NEXT:    [[DOTUPTO01:%.*]] = insertelement <2 x i32> poison, i32 [[DOTI0]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i32> [[DOTUPTO01]], i32 [[DOTI1]], i32 1
+; CHECK-NEXT:    ret <2 x i32> [[TMP1]]
+;
+  %1 = extractelement <1 x <2 x i32>*> %src, i32 %index
+  %2 = load <2 x i32>, <2 x i32>* %1, align 4
+  ret <2 x i32> %2
+}
+
+define <2 x i32> @f2b(<1 x <2 x i32>*> %src) {
+; CHECK-LABEL: @f2b(
+; CHECK-NEXT:    [[SRC_I0:%.*]] = extractelement <1 x <2 x i32>*> [[SRC:%.*]], i32 0
+; CHECK-NEXT:    [[SRC_I0_I0:%.*]] = bitcast <2 x i32>* [[SRC_I0]] to i32*
+; CHECK-NEXT:    [[SRC_I0_I1:%.*]] = getelementptr i32, i32* [[SRC_I0_I0]], i32 1
+; CHECK-NEXT:    [[DOTI0:%.*]] = load i32, i32* [[SRC_I0_I0]], align 4
+; CHECK-NEXT:    [[DOTI1:%.*]] = load i32, i32* [[SRC_I0_I1]], align 4
+; CHECK-NEXT:    [[DOTUPTO0:%.*]] = insertelement <2 x i32> poison, i32 [[DOTI0]], i32 0
+; CHECK-NEXT:    [[TMP1:%.*]] = insertelement <2 x i32> [[DOTUPTO0]], i32 [[DOTI1]], i32 1
+; CHECK-NEXT:    ret <2 x i32> [[TMP1]]
+;
+  %1 = extractelement <1 x <2 x i32>*> %src, i32 0
+  %2 = load <2 x i32>, <2 x i32>* %1, align 4
+  ret <2 x i32> %2
+}
+
+define void @f3(<1 x <2 x i32>*> %src, i32 %index, <2 x i32> %val) {
+; CHECK-LABEL: @f3(
+; CHECK-NEXT:    [[VAL_I0:%.*]] = extractelement <2 x i32> [[VAL:%.*]], i32 0
+; CHECK-NEXT:    [[VAL_I1:%.*]] = extractelement <2 x i32> [[VAL]], i32 1
+; CHECK-NEXT:    [[INDEX_IS_0:%.*]] = icmp eq i32 [[INDEX:%.*]], 0
+; CHECK-NEXT:    [[SRC_I0:%.*]] = extractelement <1 x <2 x i32>*> [[SRC:%.*]], i32 0
+; CHECK-NEXT:    [[DOTUPTO0:%.*]] = select i1 [[INDEX_IS_0]], <2 x i32>* [[SRC_I0]], <2 x i32>* undef
+; CHECK-NEXT:    [[DOTUPTO0_I0:%.*]] = bitcast <2 x i32>* [[DOTUPTO0]] to i32*
+; CHECK-NEXT:    [[DOTUPTO0_I1:%.*]] = getelementptr i32, i32* [[DOTUPTO0_I0]], i32 1
+; CHECK-NEXT:    store i32 [[VAL_I0]], i32* [[DOTUPTO0_I0]], align 4
+; CHECK-NEXT:    store i32 [[VAL_I1]], i32* [[DOTUPTO0_I1]], align 4
+; CHECK-NEXT:    ret void
+;
+  %1 = extractelement <1 x <2 x i32>*> %src, i32 %index
+  store <2 x i32> %val, <2 x i32>* %1, align 4
+  ret void
+}
+
+define void @f3b(<1 x <2 x i32>*> %src, <2 x i32> %val) {
+; CHECK-LABEL: @f3b(
+; CHECK-NEXT:    [[VAL_I0:%.*]] = extractelement <2 x i32> [[VAL:%.*]], i32 0
+; CHECK-NEXT:    [[VAL_I1:%.*]] = extractelement <2 x i32> [[VAL]], i32 1
+; CHECK-NEXT:    [[SRC_I0:%.*]] = extractelement <1 x <2 x i32>*> [[SRC:%.*]], i32 0
+; CHECK-NEXT:    [[SRC_I0_I0:%.*]] = bitcast <2 x i32>* [[SRC_I0]] to i32*
+; CHECK-NEXT:    [[SRC_I0_I1:%.*]] = getelementptr i32, i32* [[SRC_I0_I0]], i32 1
+; CHECK-NEXT:    store i32 [[VAL_I0]], i32* [[SRC_I0_I0]], align 4
+; CHECK-NEXT:    store i32 [[VAL_I1]], i32* [[SRC_I0_I1]], align 4
+; CHECK-NEXT:    ret void
+;
+  %1 = extractelement <1 x <2 x i32>*> %src, i32 0
+  store <2 x i32> %val, <2 x i32>* %1, align 4
+  ret void
+}