[InstCombine] Handle select inst when eliminating constant memcpy

[gandhi21299]

Summary

Allow iterating through SelectInst use of the alloca when
checking if it is only ever overwritten from constant memory.
Recursively determine if the SelectInst is replacable and insert
it into the Worklist if so. Finally, define a new SelectInst to
replace the old one, with both of it's values replaced according
to the WorkMap.

[gandhi21299]

• removed #include "llvm/IR/Instruction.h"
• ran update_test_checks.py on the test

[gandhi21299]

ping

[arsenm]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

327

IC.InsertNewInstWith should take care of debug info, but you're dropping any other instruction metadata that was attached

[gandhi21299]

• Rebased against the latest revision.
• Check if the true and false values of SelectInst are already in the Worklist before inserting the SelectInst itself into the Worklist.
• Define a new SelectInst while preserving the attached MetaData.
• Revise tests

[gandhi21299]

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

375 ↗

(On Diff #490302)

Need to be erased since it's never used.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

64–65

Handling for SelectInst should use same logic as PHINode (with IsOffset = true). Can you please add a test with a memcpy on the select result, which should not get transformed?

[gandhi21299]

• isOffset = true for SelectInst in isOnlyCopiedFromConstantMemory(..)
• Added a test where memcpy is performed upon the result of a SelectInst
• Removed the intrinsic in the test which is never used.

[nikic]

This is missing a test with different address space and removed alloca, but otherwise looks fine.

[gandhi21299]

In D136524#4065295, @nikic wrote:

This is missing a test with different address space and removed alloca, but otherwise looks fine.

Is there one? I can't think of a test case where the source and dest of a memcpy are in different address spaces and the alloca is removed.

[nikic]

In D136524#4065545, @gandhi21299 wrote:

In D136524#4065295, @nikic wrote:

This is missing a test with different address space and removed alloca, but otherwise looks fine.

Is there one? I can't think of a test case where the source and dest of a memcpy are in different address spaces and the alloca is removed.

A select between alloca and gep of alloca should work.

[gandhi21299]

• Added a test with varying addrspaces across the source and dest of a memcpy and the alloca is eliminated.

[nikic]

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

349 ↗

(On Diff #490509)

All of these 256 should be 32 instead.

397–399 ↗

(On Diff #490509)

Otherwise the select is simply unused, and this doesn't show the desired transform.

[gandhi21299]

llvm/test/Transforms/InstCombine/ptr-replace-alloca.ll

397–399 ↗

(On Diff #490509)

Despite the change from 256 to 32 and your comment above, I stil get the same result. Some other instcombine optimization optimizes this function out apparently.

[gandhi21299]

• Copy 32 bytes instead of 256 bytes.

[nikic]

Can you please rebase over https://github.com/llvm/llvm-project/commit/f61ee94470439c54499d43f0826cee04d99d5e9b? I think this should work. The problem was that we convert load of select into select of loads, so the select is outside the transform. Adding an extra GEP afterwards avoid this.

[gandhi21299]

Rebased

[nikic]

Something went wrong with the rebase. This should be the ptr-replace-alloca.ll from main with update_test_checks rerun.

[gandhi21299]

Inherit changes which were unintentionally eliminated during rebase.

[gandhi21299]

Run update_test_checks on ptr-replace-alloca.ll

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

267

Okay, I understand now why the new test case does not fold. Please switch this to the isAvailable() method added in https://github.com/llvm/llvm-project/commit/ed9806363beb1caf5265e75a64de904b60218093 and regenerate the test. I think then everything should work as expected.

[gandhi21299]

• use isAvailable as requested by reviewer
• update test accordingly

[gandhi21299]

Refreshing patch

[nikic]

LGTM

[gandhi21299]

Thank you for the review Nikita, I will commit this patch as soon as the pre-merge checks complete.

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

 #include "InstCombineInternal.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/Loads.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/DebugInfoMetadata.h"
+#include "llvm/IR/Instruction.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Transforms/InstCombine/InstCombiner.h"
 #include "llvm/Transforms/Utils/Local.h"
 using namespace llvm;
 using namespace PatternMatch;
 
for (auto &U : ValuePair.first->uses()) {
       auto *I = cast<Instruction>(U.getUser());
 
       if (auto *LI = dyn_cast<LoadInst>(I)) {
         // Ignore non-volatile loads, they are always ok.
         if (!LI->isSimple()) return false;
         continue;
       }
 
-      if (isa<BitCastInst>(I) || isa<AddrSpaceCastInst>(I)) {
+      if (isa<BitCastInst>(I) || isa<AddrSpaceCastInst>(I) ||
+          isa<SelectInst>(I)) {
         // If uses of the bitcast are ok, we are ok.
         ValuesToInspect.emplace_back(I, IsOffset);
         continue;
       }
       if (auto *GEP = dyn_cast<GetElementPtrInst>(I)) {
         // If the GEP has all zero indices, it doesn't offset the pointer. If it
         // doesn't, it does.
         ValuesToInspect.emplace_back(I, IsOffset || !GEP->hasAllZeroIndices());

 
 bool PointerReplacer::collectUsers(Instruction &I) {
   for (auto *U : I.users()) {
     auto *Inst = cast<Instruction>(&*U);
     if (auto *Load = dyn_cast<LoadInst>(Inst)) {
       if (Load->isVolatile())
         return false;
       Worklist.insert(Load);
+    } else if (auto *SI = dyn_cast<SelectInst>(Inst)) {
+      Worklist.insert(cast<Instruction>(SI->getTrueValue()));
+      Worklist.insert(cast<Instruction>(SI->getFalseValue()));
+      Worklist.insert(Inst);
+      if (!collectUsers(*Inst))
+        return false;
     } else if (isa<GetElementPtrInst>(Inst) || isa<BitCastInst>(Inst)) {
       Worklist.insert(Inst);
       if (!collectUsers(*Inst))
         return false;
     } else if (auto *MI = dyn_cast<MemTransferInst>(Inst)) {
       if (MI->isVolatile())
         return false;
       Worklist.insert(Inst);
     } else if (Inst->isLifetimeStartOrEnd()) {
if (auto *LT = dyn_cast<LoadInst>(I)) {
     assert(V && "Operand not replaced");
     auto *NewT = PointerType::getWithSamePointeeType(
         cast<PointerType>(BC->getType()),
         V->getType()->getPointerAddressSpace());
     auto *NewI = new BitCastInst(V, NewT);
     IC.InsertNewInstWith(NewI, *BC);
     NewI->takeName(BC);
     WorkMap[BC] = NewI;
+  } else if (auto *SI = dyn_cast<SelectInst>(I)) {
+    auto *NewSI = SelectInst::Create(SI->getCondition(),
+                                     getReplacement(SI->getTrueValue()),
+                                     getReplacement(SI->getFalseValue()));
+    IC.InsertNewInstWith(NewSI, *SI);
+    NewSI->takeName(SI);
+    WorkMap[SI] = NewSI;
   } else if (auto *MemCpy = dyn_cast<MemTransferInst>(I)) {
     auto *SrcV = getReplacement(MemCpy->getRawSource());
     // The pointer may appear in the destination of a copy, but we don't want to
     // replace it.
     if (!SrcV) {
       assert(getReplacement(MemCpy->getRawDest()) &&
              "destination not in replace list");
       return;

llvm/test/Transforms/InstCombine/replace-alloca-merge.ll

+; RUN: opt -passes=instcombine -S -o - %s | FileCheck %s
+
+target triple="amdgcn-amd-amdhsa"
+
+define i8 @remove_alloca_use_arg([32 x i8] addrspace(4)* noalias readonly align 4 dereferenceable(32) %arg, i1 %cond) {
+; CHECK-LABEL: @remove_alloca_use_arg
+; CHECK-NEXT: entry:
+; CHECK-NOT: [[ALLOCA:%.*]] = alloca [32 x i8], align 4, addrspace(5)
+; CHECK-NEXT: [[VAL1_PTR:%.*]] = getelementptr [32 x i8], [32 x i8] addrspace(4)* %arg, i64 0, i64 1
+; CHECK-NEXT: [[VAL2_PTR:%.*]] = getelementptr [32 x i8], [32 x i8] addrspace(4)* %arg, i64 0, i64 0
+; CHECK-NEXT: [[VAL1:%.*]] = load i8, i8 addrspace(4)* [[VAL1_PTR]], align 1
+; CHECK-NEXT: [[VAL2:%.*]] = load i8, i8 addrspace(4)* [[VAL2_PTR]], align 4
+; CHECK-NEXT: [[LOAD_VAL:%.*]] = select i1 %cond, i8 [[VAL1]], i8 [[VAL2]]
+; CHECK-NEXT: ret i8 [[LOAD_VAL]]
+entry:
+  %alloca = alloca [32 x i8], align 4, addrspace(5)
+  call void @llvm.memcpy.p5i8.p4i8.i64([32 x i8] addrspace(5)* %alloca, [32 x i8] addrspace(4)* %arg, i64 256, i1 false)
+  %val1 = getelementptr inbounds [32 x i8], [32 x i8] addrspace(5)* %alloca, i32 0, i32 1
+  %val2 = getelementptr inbounds [32 x i8], [32 x i8] addrspace(5)* %alloca, i32 0, i32 0
+  %ptr = select i1 %cond, i8 addrspace(5)* %val1, i8 addrspace(5)* %val2
+  %load = load i8, i8 addrspace(5)* %ptr
+  ret i8 %load
+}
+
+define i8 @volatile_load_keep_alloca([32 x i8] addrspace(4)* noalias readonly align 4 dereferenceable(32) %arg, i1 %cond) {
+; CHECK-LABEL: @volatile_load_keep_alloca
+; CHECK-NEXT: entry:
+; CHECK-NEXT: [[ALLOCA:%.*]] = alloca [32 x i8], align 4, addrspace(5)
+; CHECK-NEXT: call void @llvm.memcpy.p5a32i8.p4a32i8.i64([32 x i8] addrspace(5)* noundef nonnull align 4 dereferenceable(256) [[ALLOCA]], [32 x i8] addrspace(4)* noundef align 4 dereferenceable(256) %arg, i64 256, i1 false)
+; CHECK-NEXT: [[VAL1_PTR:%.*]] = getelementptr inbounds [32 x i8], [32 x i8] addrspace(5)* [[ALLOCA]], i64 0, i64 1
+; CHECK-NEXT: [[VAL2_PTR:%.*]] = getelementptr inbounds [32 x i8], [32 x i8] addrspace(5)* [[ALLOCA]], i64 0, i64 0
+; CHECK-NEXT: [[PTR:%.*]] = select i1 %cond, i8 addrspace(5)* [[VAL1_PTR]], i8 addrspace(5)* [[VAL2_PTR]]
+; CHECK-NEXT: [[LOAD_VAL:%.*]] = load volatile i8, i8 addrspace(5)* [[PTR]]
+; CHECK-NEXT: ret i8 [[LOAD_VAL]]
+entry:
+  %alloca = alloca [32 x i8], align 4, addrspace(5)
+  call void @llvm.memcpy.p5i8.p4i8.i64([32 x i8] addrspace(5)* %alloca, [32 x i8] addrspace(4)* %arg, i64 256, i1 false)
+  %val1 = getelementptr inbounds [32 x i8], [32 x i8] addrspace(5)* %alloca, i32 0, i32 1
+  %val2 = getelementptr inbounds [32 x i8], [32 x i8] addrspace(5)* %alloca, i32 0, i32 0
+  %ptr = select i1 %cond, i8 addrspace(5)* %val1, i8 addrspace(5)* %val2
+  %load = load volatile i8, i8 addrspace(5)* %ptr
+  ret i8 %load
+}
+
+declare void @llvm.memcpy.p5i8.p4i8.i64([32 x i8] addrspace(5)*, [32 x i8] addrspace(4)*, i64, i1)