[llvm] 9a41a80 - [AMDGPU] Handle object size and bail if assume-like intrinsic is used in PromoteAllocaToVector (#68744)

[via llvm-commits]

Author: Mariusz Sikora
Date: 2023-12-20T07:47:49+01:00
New Revision: 9a41a80e76e7b968cbec443faa2ad25384f01d6c

URL: https://github.com/llvm/llvm-project/commit/9a41a80e76e7b968cbec443faa2ad25384f01d6c
DIFF: https://github.com/llvm/llvm-project/commit/9a41a80e76e7b968cbec443faa2ad25384f01d6c.diff

LOG: [AMDGPU] Handle object size and bail if assume-like intrinsic is used in PromoteAllocaToVector (#68744)

Attached test will cause crash without this change.

We should not remove isAssumeLikeIntrinsic instruction if it is used by
other instruction.

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
    llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp b/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
index 1bed516fb5c7f8..5e73411cae9b70 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
@@ -530,6 +530,15 @@ static Value *promoteAllocaUserToVector(
       return Builder.CreateVectorSplat(VectorTy->getElementCount(), Elt);
     }
 
+    if (auto *Intr = dyn_cast<IntrinsicInst>(Inst)) {
+      if (Intr->getIntrinsicID() == Intrinsic::objectsize) {
+        Intr->replaceAllUsesWith(
+            Builder.getIntN(Intr->getType()->getIntegerBitWidth(),
+                            DL.getTypeAllocSize(VectorTy)));
+        return nullptr;
+      }
+    }
+
     llvm_unreachable("Unsupported call when promoting alloca to vector");
   }
 
@@ -773,8 +782,17 @@ bool AMDGPUPromoteAllocaImpl::tryPromoteAllocaToVector(AllocaInst &Alloca) {
       continue;
     }
 
+    if (auto *Intr = dyn_cast<IntrinsicInst>(Inst)) {
+      if (Intr->getIntrinsicID() == Intrinsic::objectsize) {
+        WorkList.push_back(Inst);
+        continue;
+      }
+    }
+
     // Ignore assume-like intrinsics and comparisons used in assumes.
     if (isAssumeLikeIntrinsic(Inst)) {
+      if (!Inst->use_empty())
+        return RejectUser(Inst, "assume-like intrinsic cannot have any users");
       UsersToRemove.push_back(Inst);
       continue;
     }

diff  --git a/llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll b/llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll
index 0bba1bdce95655..aabd5df9568370 100644
--- a/llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll
+++ b/llvm/test/CodeGen/AMDGPU/promote-alloca-mem-intrinsics.ll
@@ -53,6 +53,14 @@ define amdgpu_kernel void @promote_with_objectsize(ptr addrspace(1) %out) #0 {
   ret void
 }
 
+; CHECK-LABEL: @promote_with_objectsize_8(
+; CHECK: store i32 32, ptr addrspace(1) %out, align 4
+define amdgpu_kernel void @promote_with_objectsize_8(ptr addrspace(1) %out) #0 {
+  %alloca = alloca [8 x i32], align 4, addrspace(5)
+  %size = call i32 @llvm.objectsize.i32.p5(ptr addrspace(5) %alloca, i1 false, i1 false, i1 false)
+  store i32 %size, ptr addrspace(1) %out
+  ret void
+}
 ; CHECK-LABEL: @promote_alloca_used_twice_in_memcpy(
 ; CHECK: call void @llvm.memcpy.p3.p3.i64(ptr addrspace(3) align 8 dereferenceable(16) %arrayidx1, ptr addrspace(3) align 8 dereferenceable(16) %arrayidx2, i64 16, i1 false)
 define amdgpu_kernel void @promote_alloca_used_twice_in_memcpy(i32 %c) {