[AMDGPU] Handle invariant marks in `AMDGPUPromoteAllocaPass` #124607

shiltian · 2025-01-27T18:40:42Z

Fixes SWDEV-509327.

shiltian · 2025-01-27T18:41:01Z

[AMDGPU] Handle invariant marks in AMDGPUPromoteAllocaPass #124607 👈 (View in Graphite)
main

This stack of pull requests is managed by Graphite. Learn more about stacking.

llvmbot · 2025-01-27T18:41:22Z

@llvm/pr-subscribers-backend-amdgpu

Author: Shilei Tian (shiltian)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/124607.diff

2 Files Affected:

(modified) llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp (+26-4)
(added) llvm/test/CodeGen/AMDGPU/promote-alloca-invariant-marks.ll (+75)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp b/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
index 907f82ed7fc528..ee812cf4a81a02 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp
@@ -1556,12 +1556,34 @@ bool AMDGPUPromoteAllocaImpl::tryPromoteAllocaToLDS(AllocaInst &I,
     case Intrinsic::invariant_start:
     case Intrinsic::invariant_end:
     case Intrinsic::launder_invariant_group:
-    case Intrinsic::strip_invariant_group:
+    case Intrinsic::strip_invariant_group: {
+      SmallVector<Type *> ArgTy;
+      SmallVector<Value *> Args;
+      if (Intr->getIntrinsicID() == Intrinsic::invariant_start) {
+        Value *Size = Intr->getArgOperand(0);
+        ArgTy.emplace_back(Offset->getType());
+        Args.emplace_back(Size);
+        Args.emplace_back(Offset);
+      } else if (Intr->getIntrinsicID() == Intrinsic::invariant_end) {
+        Value *InvarientPtr = Intr->getArgOperand(0);
+        Value *Size = Intr->getArgOperand(1);
+        ArgTy.emplace_back(Offset->getType());
+        Args.emplace_back(InvarientPtr);
+        Args.emplace_back(Size);
+        Args.emplace_back(Offset);
+      } else {
+        ArgTy.emplace_back(Offset->getType());
+        Args.emplace_back(Offset);
+      }
+      Function *F = Intrinsic::getOrInsertDeclaration(
+          Intr->getModule(), Intr->getIntrinsicID(), ArgTy);
+      CallInst *NewIntr =
+          CallInst::Create(F, Args, Intr->getName(), Intr->getIterator());
+      Intr->mutateType(NewIntr->getType());
+      Intr->replaceAllUsesWith(NewIntr);
       Intr->eraseFromParent();
-      // FIXME: I think the invariant marker should still theoretically apply,
-      // but the intrinsics need to be changed to accept pointers with any
-      // address space.
       continue;
+    }
     case Intrinsic::objectsize: {
       Value *Src = Intr->getOperand(0);
 
diff --git a/llvm/test/CodeGen/AMDGPU/promote-alloca-invariant-marks.ll b/llvm/test/CodeGen/AMDGPU/promote-alloca-invariant-marks.ll
new file mode 100644
index 00000000000000..fca4be5e76daea
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/promote-alloca-invariant-marks.ll
@@ -0,0 +1,75 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -passes=amdgpu-promote-alloca %s -o - | FileCheck %s
+
+declare ptr @llvm.invariant.start.p5(i64, ptr addrspace(5) nocapture)
+declare void @llvm.invariant.end.p5(ptr, i64, ptr addrspace(5) nocapture)
+declare ptr addrspace(5) @llvm.launder.invariant.group.p5(ptr addrspace(5))
+declare ptr addrspace(5) @llvm.strip.invariant.group.p5(ptr addrspace(5))
+
+define amdgpu_kernel void @use_invariant_start_and_end() {
+; CHECK-LABEL: define amdgpu_kernel void @use_invariant_start_and_end() {
+; CHECK-NEXT:  [[BB:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call noalias nonnull dereferenceable(64) ptr addrspace(4) @llvm.amdgcn.dispatch.ptr()
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, ptr addrspace(4) [[TMP0]], i64 1
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr addrspace(4) [[TMP1]], align 4, !invariant.load [[META0:![0-9]+]]
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds i32, ptr addrspace(4) [[TMP0]], i64 2
+; CHECK-NEXT:    [[TMP4:%.*]] = load i32, ptr addrspace(4) [[TMP3]], align 4, !range [[RNG1:![0-9]+]], !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP5:%.*]] = lshr i32 [[TMP2]], 16
+; CHECK-NEXT:    [[TMP6:%.*]] = call range(i32 0, 1024) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP7:%.*]] = call range(i32 0, 1024) i32 @llvm.amdgcn.workitem.id.y()
+; CHECK-NEXT:    [[TMP8:%.*]] = call range(i32 0, 1024) i32 @llvm.amdgcn.workitem.id.z()
+; CHECK-NEXT:    [[TMP9:%.*]] = mul nuw nsw i32 [[TMP5]], [[TMP4]]
+; CHECK-NEXT:    [[TMP10:%.*]] = mul i32 [[TMP9]], [[TMP6]]
+; CHECK-NEXT:    [[TMP11:%.*]] = mul nuw nsw i32 [[TMP7]], [[TMP4]]
+; CHECK-NEXT:    [[TMP12:%.*]] = add i32 [[TMP10]], [[TMP11]]
+; CHECK-NEXT:    [[TMP13:%.*]] = add i32 [[TMP12]], [[TMP8]]
+; CHECK-NEXT:    [[TMP14:%.*]] = getelementptr inbounds [1024 x i32], ptr addrspace(3) @use_invariant_start_and_end.alloca, i32 0, i32 [[TMP13]]
+; CHECK-NEXT:    [[INVARIENT1:%.*]] = call ptr @llvm.invariant.start.p3(i64 0, ptr addrspace(3) [[TMP14]])
+; CHECK-NEXT:    store <2 x i1> zeroinitializer, ptr [[INVARIENT1]], align 1
+; CHECK-NEXT:    call void @llvm.invariant.end.p3(ptr [[INVARIENT1]], i64 0, ptr addrspace(3) [[TMP14]])
+; CHECK-NEXT:    ret void
+;
+bb:
+  %alloca = alloca i32, align 4, addrspace(5)
+  %invarient = call ptr @llvm.invariant.start.p5(i64 0, ptr addrspace(5) %alloca)
+  store <2 x i1> zeroinitializer, ptr %invarient, align 1
+  call void @llvm.invariant.end.p5(ptr %invarient, i64 0, ptr addrspace(5) %alloca)
+  ret void
+}
+
+define amdgpu_kernel void @use_invariant_group_and_strip() {
+; CHECK-LABEL: define amdgpu_kernel void @use_invariant_group_and_strip() {
+; CHECK-NEXT:  [[BB:.*:]]
+; CHECK-NEXT:    [[TMP0:%.*]] = call noalias nonnull dereferenceable(64) ptr addrspace(4) @llvm.amdgcn.dispatch.ptr()
+; CHECK-NEXT:    [[TMP1:%.*]] = getelementptr inbounds i32, ptr addrspace(4) [[TMP0]], i64 1
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, ptr addrspace(4) [[TMP1]], align 4, !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr inbounds i32, ptr addrspace(4) [[TMP0]], i64 2
+; CHECK-NEXT:    [[TMP4:%.*]] = load i32, ptr addrspace(4) [[TMP3]], align 4, !range [[RNG1]], !invariant.load [[META0]]
+; CHECK-NEXT:    [[TMP5:%.*]] = lshr i32 [[TMP2]], 16
+; CHECK-NEXT:    [[TMP6:%.*]] = call range(i32 0, 1024) i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NEXT:    [[TMP7:%.*]] = call range(i32 0, 1024) i32 @llvm.amdgcn.workitem.id.y()
+; CHECK-NEXT:    [[TMP8:%.*]] = call range(i32 0, 1024) i32 @llvm.amdgcn.workitem.id.z()
+; CHECK-NEXT:    [[TMP9:%.*]] = mul nuw nsw i32 [[TMP5]], [[TMP4]]
+; CHECK-NEXT:    [[TMP10:%.*]] = mul i32 [[TMP9]], [[TMP6]]
+; CHECK-NEXT:    [[TMP11:%.*]] = mul nuw nsw i32 [[TMP7]], [[TMP4]]
+; CHECK-NEXT:    [[TMP12:%.*]] = add i32 [[TMP10]], [[TMP11]]
+; CHECK-NEXT:    [[TMP13:%.*]] = add i32 [[TMP12]], [[TMP8]]
+; CHECK-NEXT:    [[TMP14:%.*]] = getelementptr inbounds [1024 x i32], ptr addrspace(3) @use_invariant_group_and_strip.alloca, i32 0, i32 [[TMP13]]
+; CHECK-NEXT:    [[INVARIENT2:%.*]] = call ptr addrspace(3) @llvm.launder.invariant.group.p3(ptr addrspace(3) [[TMP14]])
+; CHECK-NEXT:    store <2 x i1> zeroinitializer, ptr addrspace(3) [[INVARIENT2]], align 1
+; CHECK-NEXT:    [[STRIP1:%.*]] = call ptr addrspace(3) @llvm.strip.invariant.group.p3(ptr addrspace(3) [[TMP14]])
+; CHECK-NEXT:    store <2 x i1> zeroinitializer, ptr addrspace(3) [[STRIP1]], align 1
+; CHECK-NEXT:    ret void
+;
+bb:
+  %alloca = alloca i32, align 4, addrspace(5)
+  %invarient = call ptr addrspace(5) @llvm.launder.invariant.group.p5(ptr addrspace(5) %alloca)
+  store <2 x i1> zeroinitializer, ptr addrspace(5) %invarient, align 1
+  %strip = call ptr addrspace(5) @llvm.strip.invariant.group.p5(ptr addrspace(5) %alloca)
+  store <2 x i1> zeroinitializer, ptr addrspace(5) %strip, align 1
+  ret void
+}
+;.
+; CHECK: [[META0]] = !{}
+; CHECK: [[RNG1]] = !{i32 0, i32 1025}
+;.

rampitec

LGTM

arsenm · 2025-01-28T01:22:40Z

llvm/test/CodeGen/AMDGPU/promote-alloca-invariant-marks.ll

+;
+bb:
+  %alloca = alloca i32, align 4, addrspace(5)
+  %invarient = call ptr @llvm.invariant.start.p5(i64 0, ptr addrspace(5) %alloca)


Typo invarient

oops, my poor English…

arsenm · 2025-01-28T01:23:26Z

llvm/test/CodeGen/AMDGPU/promote-alloca-invariant-marks.ll

+  %alloca = alloca i32, align 4, addrspace(5)
+  %invarient = call ptr @llvm.invariant.start.p5(i64 0, ptr addrspace(5) %alloca)
+  store <2 x i1> zeroinitializer, ptr %invarient, align 1
+  call void @llvm.invariant.end.p5(ptr %invarient, i64 0, ptr addrspace(5) %alloca)


Why is the return type an addrspace(0) pointer? Should this really be using a token?

Why is the return type an addrspace(0) pointer?

That is based on the LangRef.

Should this really be using a token?

I'm not sure what you meant by using a token.

I mean the token IR type. What is this pointer value used for

I'm not very sure actually, but based on llvm/test/Transforms/RewriteStatepointsForGC/drop-invalid-metadata.ll, it seems feasible.

%invst = call ptr @llvm.invariant.start.p1(i64 1, ptr addrspace(1) %v1) ... %foo = call i32 @escaping.invariant.start(ptr %invst)

arsenm · 2025-01-28T01:23:42Z

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

+        ArgTy.emplace_back(Offset->getType());
+        Args.emplace_back(Size);
+        Args.emplace_back(Offset);


Assign from initializer list?

Sure, will do a follow up.

arsenm · 2025-01-28T01:24:12Z

llvm/lib/Target/AMDGPU/AMDGPUPromoteAlloca.cpp

+        Args.emplace_back(Size);
+        Args.emplace_back(Offset);
+      } else {
+        ArgTy.emplace_back(Offset->getType());


This is common to all cases, you could initialize the vectors with it

Sure, will do a follow up.

llvmbot added the backend:AMDGPU label Jan 27, 2025

shiltian requested review from jayfoad, arsenm, rampitec and Pierre-vh January 27, 2025 18:41

rampitec approved these changes Jan 27, 2025

View reviewed changes

[AMDGPU] Handle invariant marks in AMDGPUPromoteAllocaPass

5e48dd7

shiltian force-pushed the users/shiltian/handle-invariant-marks branch from 5f17deb to 5e48dd7 Compare January 27, 2025 22:30

shiltian merged commit 3b2b7ec into main Jan 27, 2025
4 of 6 checks passed

shiltian deleted the users/shiltian/handle-invariant-marks branch January 27, 2025 22:30

arsenm reviewed Jan 28, 2025

View reviewed changes

shiltian added a commit that referenced this pull request Jan 28, 2025

[NFC][AMDGPU] Improve code introduced in #124607

cde3d5a

shiltian added a commit that referenced this pull request Jan 28, 2025

[NFC][AMDGPU] Improve code introduced in #124607 (#124672)

6e41055

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[AMDGPU] Handle invariant marks in `AMDGPUPromoteAllocaPass` #124607

[AMDGPU] Handle invariant marks in `AMDGPUPromoteAllocaPass` #124607

Uh oh!

shiltian commented Jan 27, 2025 •

edited

Loading

Uh oh!

shiltian commented Jan 27, 2025

Uh oh!

llvmbot commented Jan 27, 2025

Uh oh!

rampitec left a comment

Uh oh!

Uh oh!

arsenm Jan 28, 2025

Uh oh!

shiltian Jan 28, 2025

Uh oh!

arsenm Jan 28, 2025

Uh oh!

shiltian Jan 28, 2025

Uh oh!

arsenm Jan 28, 2025

Uh oh!

shiltian Jan 28, 2025

Uh oh!

arsenm Jan 28, 2025

Uh oh!

shiltian Jan 28, 2025

Uh oh!

arsenm Jan 28, 2025

Uh oh!

shiltian Jan 28, 2025

Uh oh!

Uh oh!

[AMDGPU] Handle invariant marks in AMDGPUPromoteAllocaPass #124607

[AMDGPU] Handle invariant marks in AMDGPUPromoteAllocaPass #124607

Uh oh!

Conversation

shiltian commented Jan 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

shiltian commented Jan 27, 2025

Uh oh!

llvmbot commented Jan 27, 2025

Uh oh!

rampitec left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

[AMDGPU] Handle invariant marks in `AMDGPUPromoteAllocaPass` #124607

[AMDGPU] Handle invariant marks in `AMDGPUPromoteAllocaPass` #124607

shiltian commented Jan 27, 2025 •

edited

Loading