[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

399

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

392

Need to be erased since it's never used.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombineLoadStoreAlloca.cpp

82

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

All of these 256 should be 32 instead.

397–399

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

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

324

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

for (auto &U : Value->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<PHINode>(I)) {
+      if (isa<PHINode, SelectInst>(I)) {
         // We set IsOffset=true, to forbid the memcpy from occurring after the
         // phi: If one of the phi operands is not based on the alloca, we
         // would incorrectly omit a write.
         Worklist.emplace_back(I, true);
         continue;
       }
-      if (isa<BitCastInst>(I) || isa<AddrSpaceCastInst>(I)) {
+      if (isa<BitCastInst, AddrSpaceCastInst>(I)) {
         // If uses of the bitcast are ok, we are ok.
         Worklist.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.
         Worklist.emplace_back(I, IsOffset || !GEP->hasAllZeroIndices());
if (auto *Load = dyn_cast<LoadInst>(Inst)) {
           })) {
         ValuesToRevisit.insert(Inst);
         continue;
       }
 
       Worklist.insert(PHI);
       if (!collectUsersRecursive(*PHI))
         return false;
+    } else if (auto *SI = dyn_cast<SelectInst>(Inst)) {
+      if (!isa<Instruction>(SI->getTrueValue()) ||
+          !isa<Instruction>(SI->getFalseValue()))
+        return false;
+
+      if (!Worklist.contains(cast<Instruction>(SI->getTrueValue())) ||
+          !Worklist.contains(cast<Instruction>(SI->getFalseValue()))) {
+        ValuesToRevisit.insert(Inst);
+        continue;
+      }
+      Worklist.insert(SI);
+      if (!collectUsersRecursive(*SI))
+        return false;
     } else if (isa<GetElementPtrInst, BitCastInst>(Inst)) {
       Worklist.insert(Inst);
       if (!collectUsersRecursive(*Inst))
         return false;
     } else if (auto *MI = dyn_cast<MemTransferInst>(Inst)) {
       if (MI->isVolatile())
         return false;
       Worklist.insert(Inst);
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()), SI->getName(), nullptr, SI);
+    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/ptr-replace-alloca.ll

loop:
   %ptr.next = getelementptr i8, ptr %ptr, i64 4
   br i1 %c, label %exit, label %loop
 
 exit:
   %v = load i32, ptr %ptr
   ret i32 %v
 }
 
+define i8 @select_same_addrspace_remove_alloca(i1 %cond, ptr %p) {
+; CHECK-LABEL: @select_same_addrspace_remove_alloca(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[PTR:%.*]] = select i1 %cond, ptr [[G1:@.*]], ptr [[P:%.*]]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr [[PTR]], align 1
+; CHECK-NEXT:    ret i8 [[LOAD]]
+;
+entry:
+  %alloca = alloca [32 x i8]
+  call void @llvm.memcpy.p0.p0.i64(ptr %alloca, ptr @g1, i64 256, i1 false)
+  %ptr = select i1 %cond, ptr %alloca, ptr %p
+  %load = load i8, ptr %ptr
+  ret i8 %load
+}
+
+define i8 @select_after_memcpy_keep_alloca(i1 %cond, ptr %p) {
+; CHECK-LABEL: define i8 @select_after_memcpy_keep_alloca(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:   [[ALLOCA:%.*]] = alloca [32 x i8], align 1
+; CHECK-NEXT:   [[PTR:%.*]] = select i1 %cond, ptr [[ALLOCA]], ptr [[P:%.*]]
+; CHECK-NEXT:   call void @llvm.memcpy.p0.p0.i64(ptr noundef nonnull align 1 dereferenceable(256) [[PTR]], ptr noundef nonnull align 16 dereferenceable(256) [[G1:@.*]], i64 256, i1 false)
+; CHECK-NEXT:   [[LOAD:%.*]] = load i8, ptr [[PTR]], align 1
+; CHECK-NEXT:   ret i8 [[LOAD]]
+;
+entry:
+  %alloca = alloca [32 x i8]
+  %ptr = select i1 %cond, ptr %alloca, ptr %p
+  call void @llvm.memcpy.p0.p0.i64(ptr %ptr, ptr @g1, i64 256, i1 false)
+  %load = load i8, ptr %ptr
+  ret i8 %load
+}
+
+define i8 @select_diff_addrspace_keep_alloca(i1 %cond, ptr addrspace(1) %p) {
+; CHECK-LABEL: @select_diff_addrspace_keep_alloca(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    %alloca = alloca [32 x i8], align 1, addrspace(1)
+; CHECK-NEXT:    call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) noundef align 1 dereferenceable(256) %alloca, ptr noundef nonnull align 16 dereferenceable(256) @g1, i64 256, i1 false)
+; CHECK-NEXT:    %ptr = select i1 %cond, ptr addrspace(1) %alloca, ptr addrspace(1) %p
+; CHECK-NEXT:    %load = load i8, ptr addrspace(1) %ptr, align 1
+; CHECK-NEXT:    ret i8 %load
+;
+entry:
+  %alloca = alloca [32 x i8], addrspace(1)
+  call void @llvm.memcpy.p1.p0.i64(ptr addrspace(1) %alloca, ptr @g1, i64 256, i1 false)
+  %ptr = select i1 %cond, ptr addrspace(1) %alloca, ptr addrspace(1) %p
+  %load = load i8, ptr addrspace(1) %ptr
+  ret i8 %load
+}
+
 declare void @llvm.memcpy.p1.p0.i64(ptr addrspace(1), ptr, i64, i1)
 declare void @llvm.memcpy.p0.p0.i64(ptr, ptr, i64, i1)
 declare void @llvm.memcpy.p0.p1.i64(ptr, ptr addrspace(1), i64, i1)