[llvm] 758fea0 - [InferAddressSpaces] Handle llvm.lifetime (#141045)

[via llvm-commits]

Author: QiYue
Date: 2025-05-22T18:06:01+02:00
New Revision: 758fea0e995e6128022e5cd2605a92222e130837

URL: https://github.com/llvm/llvm-project/commit/758fea0e995e6128022e5cd2605a92222e130837
DIFF: https://github.com/llvm/llvm-project/commit/758fea0e995e6128022e5cd2605a92222e130837.diff

LOG: [InferAddressSpaces] Handle llvm.lifetime (#141045)

Co-authored-by: Zhenhao Yang <zhenhao.yang at nio.com>
Co-authored-by: Matt Arsenault <arsenm2 at gmail.com>

Added: 
    llvm/test/Transforms/InferAddressSpaces/AMDGPU/lifetime.ll
    llvm/test/Transforms/InferAddressSpaces/NVPTX/lifetime.ll

Modified: 
    llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index 1b7cecc7ceb6a..66836ef05d5db 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -428,6 +428,14 @@ bool InferAddressSpacesImpl::rewriteIntrinsicOperands(IntrinsicInst *II,
     II->replaceUsesOfWith(OldV, NewV);
     return true;
   }
+  case Intrinsic::lifetime_start:
+  case Intrinsic::lifetime_end: {
+    Function *NewDecl = Intrinsic::getOrInsertDeclaration(
+        M, II->getIntrinsicID(), {NewV->getType()});
+    II->setArgOperand(1, NewV);
+    II->setCalledFunction(NewDecl);
+    return true;
+  }
   default: {
     Value *Rewrite = TTI->rewriteIntrinsicWithAddressSpace(II, OldV, NewV);
     if (!Rewrite)
@@ -479,6 +487,12 @@ void InferAddressSpacesImpl::collectRewritableIntrinsicOperands(
 
     break;
   }
+  case Intrinsic::lifetime_start:
+  case Intrinsic::lifetime_end: {
+    appendsFlatAddressExpressionToPostorderStack(II->getArgOperand(1),
+                                                 PostorderStack, Visited);
+    break;
+  }
   default:
     SmallVector<int, 2> OpIndexes;
     if (TTI->collectFlatAddressOperands(OpIndexes, IID)) {

diff  --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/lifetime.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/lifetime.ll
new file mode 100644
index 0000000000000..d39a0b35ce811
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/lifetime.ll
@@ -0,0 +1,22 @@
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=infer-address-spaces %s | FileCheck %s
+
+define i32 @lifetime_flat_pointer() {
+; CHECK-LABEL: define i32 @lifetime_flat_pointer() {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, align 4, addrspace(5)
+; CHECK-NEXT:    call void @llvm.lifetime.start.p5(i64 4, ptr addrspace(5) [[ALLOCA]])
+; CHECK-NEXT:    store i32 1, ptr addrspace(5) [[ALLOCA]], align 4
+; CHECK-NEXT:    %ret = load i32, ptr addrspace(5) [[ALLOCA]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p5(i64 4, ptr addrspace(5) [[ALLOCA]])
+; CHECK-NEXT:    ret i32 %ret
+;
+  %alloca = alloca i32, align 4, addrspace(5)
+  %flat = addrspacecast ptr addrspace(5) %alloca to ptr
+  call void @llvm.lifetime.start.p0(i64 4 , ptr %flat)
+  store i32 1, ptr %flat, align 4
+  %ret = load i32, ptr %flat, align 4
+  call void @llvm.lifetime.end.p0(i64 4 , ptr %flat)
+  ret i32 %ret
+}
+
+declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
+declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)

diff  --git a/llvm/test/Transforms/InferAddressSpaces/NVPTX/lifetime.ll b/llvm/test/Transforms/InferAddressSpaces/NVPTX/lifetime.ll
new file mode 100644
index 0000000000000..8bf63127ba636
--- /dev/null
+++ b/llvm/test/Transforms/InferAddressSpaces/NVPTX/lifetime.ll
@@ -0,0 +1,27 @@
+; RUN: opt -S -passes=infer-address-spaces %s | FileCheck %s
+
+target triple = "nvptx64-nvidia-cuda"
+
+define i32 @lifetime_flat_pointer() {
+; CHECK-LABEL: define i32 @lifetime_flat_pointer() {
+; CHECK-NEXT:    [[ALLOCA:%.*]] = alloca i32, align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[ALLOCA]] to ptr addrspace(5)
+; CHECK-NEXT:    call void @llvm.lifetime.start.p5(i64 4, ptr addrspace(5) [[TMP1]])
+; CHECK-NEXT:    store i32 1, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    %ret = load i32, ptr addrspace(5) [[TMP1]], align 4
+; CHECK-NEXT:    call void @llvm.lifetime.end.p5(i64 4, ptr addrspace(5) [[TMP1]])
+; CHECK-NEXT:    ret i32 %ret
+;
+  %alloca = alloca i32, align 4
+  %1 = addrspacecast ptr %alloca to ptr addrspace(5)
+  %2 = addrspacecast ptr addrspace(5) %1 to ptr
+  %3 = addrspacecast ptr addrspace(5) %1 to ptr
+  call void @llvm.lifetime.start.p0(i64 4, ptr %2)
+  store i32 1, ptr addrspace(5) %1, align 4
+  %ret = load i32, ptr addrspace(5) %1, align 4
+  call void @llvm.lifetime.end.p0(i64 4, ptr %3)
+  ret i32 %ret
+}
+
+declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture)
+declare void @llvm.lifetime.end.p0(i64 immarg, ptr nocapture)