[DSE][MemLoc] Handle intrinsics more generically

[nikic]

Summary

Remove the special casing for intrinsics in MemoryLocation::getForDest() and handle them through the general attribute based code. On the DSE side, this means that isRemovable() now needs to handle more than a hardcoded list of intrinsics. We consider everything apart from volatile accesses memory intrinsics and lifetime markers to be removable.

This allows us to perform DSE on intrinsics that DSE has not been specially been taught about, using a matrix store as an example here.

There is an interesting test change for invariant.start, but I believe that optimization is correct. It only looks a bit odd because the code is immediate UB anyway.

[nikic]

Further simplify code.

[fhahn]

LGTM, thanks!

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

185–193

might be worth an assert here

[nikic]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

185–193

In some cases we call isRemovable() before getLocForWriteEx() (but still bail out if either fails). I could add an assert after reordering the calls.

[fhahn]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

185–193

That might be a good way to prevent accidental mis-use in future.

[nikic]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

185–193

Assert added in https://github.com/llvm/llvm-project/commit/034e66e76c6f0f76aa7c67ba4edf4eba108bcf51, though it took a few more preliminary changes than anticipated.

[fhahn]

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

185–193

Great, thanks!

llvm/lib/Analysis/MemoryLocation.cpp

 
 MemoryLocation MemoryLocation::getForDest(const AnyMemIntrinsic *MI) {
   assert(MI->getRawDest() == MI->getArgOperand(0));
   return getForArgument(MI, 0, nullptr);
 }
 
 Optional<MemoryLocation>
 MemoryLocation::getForDest(const CallBase *CB, const TargetLibraryInfo &TLI) {
-  if (const IntrinsicInst *II = dyn_cast<IntrinsicInst>(CB)) {
-    if (auto *MemInst = dyn_cast<AnyMemIntrinsic>(CB))
-      return getForDest(MemInst);
-
-    switch (II->getIntrinsicID()) {
-    default:
-      return None;
-    case Intrinsic::init_trampoline:
-      return MemoryLocation::getForArgument(CB, 0, TLI);
-    case Intrinsic::masked_store:
-      return MemoryLocation::getForArgument(CB, 1, TLI);
-    }
-  }
-
   if (!CB->onlyAccessesArgMemory())
     return None;
 
   if (CB->hasOperandBundles())
     // TODO: remove implementation restriction
     return None;
 
   Value *UsedV = nullptr;

llvm/lib/Transforms/Scalar/DeadStoreElimination.cpp

 
 /// If the value of this instruction and the memory it writes to is unused, may
 /// we delete this instruction?
 static bool isRemovable(Instruction *I) {
   // Don't remove volatile/atomic stores.
   if (StoreInst *SI = dyn_cast<StoreInst>(I))
     return SI->isUnordered();
 
-  if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(I)) {
-    switch (II->getIntrinsicID()) {
-    default: llvm_unreachable("Does not have LocForWrite");
-    case Intrinsic::lifetime_end:
-      // Never remove dead lifetime_end's, e.g. because it is followed by a
-      // free.
-      return false;
-    case Intrinsic::init_trampoline:
-      // Always safe to remove init_trampoline.
-      return true;
-    case Intrinsic::memset:
-    case Intrinsic::memmove:
-    case Intrinsic::memcpy:
-    case Intrinsic::memcpy_inline:
-      // Don't remove volatile memory intrinsics.
-      return !cast<MemIntrinsic>(II)->isVolatile();
-    case Intrinsic::memcpy_element_unordered_atomic:
-    case Intrinsic::memmove_element_unordered_atomic:
-    case Intrinsic::memset_element_unordered_atomic:
-    case Intrinsic::masked_store:
-      return true;
-    }
-  }
+  // Note: only get here for calls with analyzable writes.
+  if (auto *CB = dyn_cast<CallBase>(I)) {
+    if (auto *MI = dyn_cast<MemIntrinsic>(CB))
+      return !MI->isVolatile();
+
+    // Never remove dead lifetime intrinsics, e.g. because they are followed by
+    // a free.
+    if (CB->isLifetimeStartOrEnd())
+      return false;
 
-  // note: only get here for calls with analyzable writes.
-  if (auto *CB = dyn_cast<CallBase>(I))
     return CB->use_empty() && CB->willReturn() && CB->doesNotThrow();
+  }
 
   return false;
 }
 
 /// Returns true if the end of this instruction can be safely shortened in
 /// length.
 static bool isShortenableAtTheEnd(Instruction *I) {
   // Don't shorten stores for now

llvm/test/Transforms/DeadStoreElimination/invariant.start.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
 ; Test to make sure llvm.invariant.start calls are not treated as clobbers.
 ; RUN: opt < %s -basic-aa -dse -S | FileCheck %s
 
 declare {}* @llvm.invariant.start.p0i8(i64, i8* nocapture) nounwind readonly
 
-; We cannot remove the store 1 to %p.
-; FIXME: By the semantics of invariant.start, the store 3 to p is unreachable.
+; We could remove either store here. The first store is dead in the
+; conventional sense, because there is a later killing store. The second store
+; is undefined behavior by the semantics of invariant.start, and as such
+; unreachable.
 define void @test(i8 *%p) {
 ; CHECK-LABEL: @test(
-; CHECK-NEXT:    store i8 1, i8* [[P:%.*]], align 4
-; CHECK-NEXT:    [[I:%.*]] = call {}* @llvm.invariant.start.p0i8(i64 1, i8* [[P]])
-; CHECK-NEXT:    store i8 3, i8* [[P]], align 4
+; CHECK-NEXT:    store i8 3, i8* [[P:%.*]], align 4
 ; CHECK-NEXT:    ret void
 ;
   store i8 1, i8* %p, align 4
   %i = call {}* @llvm.invariant.start.p0i8(i64 1, i8* %p)
   store i8 3, i8* %p, align 4
   ret void
 }
 

llvm/test/Transforms/DeadStoreElimination/simple.ll

 ; CHECK-NEXT:    ret void
 ;
   %storage = alloca [10 x i8], align 16    ; <[10 x i8]*> [#uses=1]
   %cast = getelementptr [10 x i8], [10 x i8]* %storage, i32 0, i32 0    ; <i8*> [#uses=1]
   call void @llvm.init.trampoline( i8* %cast, i8* bitcast (void ()* @test11f to i8*), i8* null )    ; <i8*> [#uses=1]
   ret void
 }
 
-; TODO: Specialized store intrinsics should be removed if dead.
+; Specialized store intrinsics should be removed if dead.
 define void @test_matrix_store(i64 %stride) {
 ; CHECK-LABEL: @test_matrix_store(
-; CHECK-NEXT:    [[A:%.*]] = alloca [6 x float], align 4
-; CHECK-NEXT:    [[CAST:%.*]] = bitcast [6 x float]* [[A]] to float*
-; CHECK-NEXT:    call void @llvm.matrix.column.major.store.v6f32.i64(<6 x float> zeroinitializer, float* [[CAST]], i64 [[STRIDE:%.*]], i1 false, i32 3, i32 2)
 ; CHECK-NEXT:    ret void
 ;
   %a = alloca [6 x float]
   %cast = bitcast [6 x float]* %a to float*
   call void @llvm.matrix.column.major.store(<6 x float> zeroinitializer, float* %cast, i64 %stride, i1 false, i32 3, i32 2)
   ret void
 }