[mlir][rocdl] Add GlobalLoadAsyncToLDS operation #165374

pabloantoniom · 2025-10-28T11:58:38Z

Adds global.load.async.to.lds op to rocdl, supporting b8, b32, b64 and b128. The op is lowered to the appropiate llvm.amdgcn.global.load.async.to.lds.bXX intrinsic.

This is available on gfx1250+

llvmbot · 2025-10-28T11:59:11Z

@llvm/pr-subscribers-mlir

Author: Pablo Antonio Martinez (pabloantoniom)

Changes

Adds global.load.async.to.lds op to rocdl, supporting b8, b32, b64 and b128. The op is lowered to the appropiate llvm.amdgcn.global.load.async.to.lds.bXX intrinsic.

This is available on gfx1250+

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

3 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+33)
(modified) mlir/test/Dialect/LLVMIR/rocdl.mlir (+13)
(modified) mlir/test/Target/LLVMIR/rocdl.mlir (+24)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index d2df244eb9363..3fcbbe52748f5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -663,6 +663,39 @@ def ROCDL_GlobalLoadLDSOp :
   }];
 }
 
+//===---------------------------------------------------------------------===//
+// Async load to LDS intrinsic (available in GFX1250)
+//===---------------------------------------------------------------------===//
+
+class ROCDL_GlobalLoadAsyncToLDSOp<string mnemonic> :
+  ROCDL_IntrOp<mnemonic, [], [], [], 0, 0, 1, 0, [2, 3], ["offset", "aux"]> {
+  dag args = (ins Arg<LLVM_AnyPointer, "", [MemRead]>:$globalPtr,
+                 Arg<ROCDLBufferLDS, "", [MemWrite]>:$ldsPtr,
+                 I32Attr:$offset,
+                 I32Attr:$aux);
+  let arguments = !con(args, baseArgs);
+  let assemblyFormat = [{
+    $globalPtr `,`  $ldsPtr `,` $offset `,` $aux
+    attr-dict `:` type($globalPtr)
+  }];
+  let description = [{
+    Loads data asynchronously from a global memory pointer to a local data
+    store (LDS) pointer.
+
+    Available on gfx1250+.
+  }];
+  let extraClassDefinition = [{
+    ::llvm::SmallVector<::mlir::Value> $cppClass::getAccessedOperands() {
+      return {getGlobalPtr(), getLdsPtr()};
+    }
+  }];
+}
+
+def ROCDL_GlobalLoadAsyncToLDSB8Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b8">;
+def ROCDL_GlobalLoadAsyncToLDSB32Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b32">;
+def ROCDL_GlobalLoadAsyncToLDSB64Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b64">;
+def ROCDL_GlobalLoadAsyncToLDSB128Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b128">;
+
 //===---------------------------------------------------------------------===//
 // Operations on raw buffer resources (stride of 0, bounds checks either off or in
 // raw buffer mode).
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index d270ee8b089aa..47464abd610f9 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -664,6 +664,19 @@ llvm.func @rocdl.global.load.lds(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
   llvm.return
 }
 
+llvm.func @rocdl.global.load.async.to.lds(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK-LABEL @rocdl.global.load.async.to.lds
+  // CHECK: rocdl.global.load.async.to.lds.b8 %{{.*}}, %{{.*}}, 0, 0
+  // CHECK: rocdl.global.load.async.to.lds.b32 %{{.*}}, %{{.*}}, 0, 0
+  // CHECK: rocdl.global.load.async.to.lds.b64 %{{.*}}, %{{.*}}, 0, 0
+  // CHECK: rocdl.global.load.async.to.lds.b128 %{{.*}}, %{{.*}}, 0, 0
+  rocdl.global.load.async.to.lds.b8 %src, %dst, 0, 0 : <1>
+  rocdl.global.load.async.to.lds.b32 %src, %dst, 0, 0 : <1>
+  rocdl.global.load.async.to.lds.b64 %src, %dst, 0, 0 : <1>
+  rocdl.global.load.async.to.lds.b128 %src, %dst, 0, 0 : <1>
+  llvm.return
+}
+
 llvm.func @rocdl.make.buffer.rsrc(%ptr : !llvm.ptr,
                                   %stride : i16,
                                   %numRecords : i64,
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 30126f6bff05a..5ae9f11360df4 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -1040,6 +1040,30 @@ llvm.func @rocdl.global.load.lds(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
   llvm.return
 }
 
+llvm.func @rocdl.global.load.async.lds.b8(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b8
+  rocdl.global.load.async.to.lds.b8 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @rocdl.global.load.async.lds.b32(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b32
+  rocdl.global.load.async.to.lds.b32 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @rocdl.global.load.async.lds.b64(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b64
+  rocdl.global.load.async.to.lds.b64 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @rocdl.global.load.async.lds.b128(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b128
+  rocdl.global.load.async.to.lds.b128 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
 llvm.func @rocdl.make.buffer.rsrc(%ptr : !llvm.ptr,
                                   %stride : i16,
                                   %numRecords : i64,

llvmbot · 2025-10-28T11:59:12Z

@llvm/pr-subscribers-mlir-llvm

Author: Pablo Antonio Martinez (pabloantoniom)

Changes

Adds global.load.async.to.lds op to rocdl, supporting b8, b32, b64 and b128. The op is lowered to the appropiate llvm.amdgcn.global.load.async.to.lds.bXX intrinsic.

This is available on gfx1250+

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

3 Files Affected:

(modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+33)
(modified) mlir/test/Dialect/LLVMIR/rocdl.mlir (+13)
(modified) mlir/test/Target/LLVMIR/rocdl.mlir (+24)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index d2df244eb9363..3fcbbe52748f5 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -663,6 +663,39 @@ def ROCDL_GlobalLoadLDSOp :
   }];
 }
 
+//===---------------------------------------------------------------------===//
+// Async load to LDS intrinsic (available in GFX1250)
+//===---------------------------------------------------------------------===//
+
+class ROCDL_GlobalLoadAsyncToLDSOp<string mnemonic> :
+  ROCDL_IntrOp<mnemonic, [], [], [], 0, 0, 1, 0, [2, 3], ["offset", "aux"]> {
+  dag args = (ins Arg<LLVM_AnyPointer, "", [MemRead]>:$globalPtr,
+                 Arg<ROCDLBufferLDS, "", [MemWrite]>:$ldsPtr,
+                 I32Attr:$offset,
+                 I32Attr:$aux);
+  let arguments = !con(args, baseArgs);
+  let assemblyFormat = [{
+    $globalPtr `,`  $ldsPtr `,` $offset `,` $aux
+    attr-dict `:` type($globalPtr)
+  }];
+  let description = [{
+    Loads data asynchronously from a global memory pointer to a local data
+    store (LDS) pointer.
+
+    Available on gfx1250+.
+  }];
+  let extraClassDefinition = [{
+    ::llvm::SmallVector<::mlir::Value> $cppClass::getAccessedOperands() {
+      return {getGlobalPtr(), getLdsPtr()};
+    }
+  }];
+}
+
+def ROCDL_GlobalLoadAsyncToLDSB8Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b8">;
+def ROCDL_GlobalLoadAsyncToLDSB32Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b32">;
+def ROCDL_GlobalLoadAsyncToLDSB64Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b64">;
+def ROCDL_GlobalLoadAsyncToLDSB128Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b128">;
+
 //===---------------------------------------------------------------------===//
 // Operations on raw buffer resources (stride of 0, bounds checks either off or in
 // raw buffer mode).
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index d270ee8b089aa..47464abd610f9 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -664,6 +664,19 @@ llvm.func @rocdl.global.load.lds(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
   llvm.return
 }
 
+llvm.func @rocdl.global.load.async.to.lds(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK-LABEL @rocdl.global.load.async.to.lds
+  // CHECK: rocdl.global.load.async.to.lds.b8 %{{.*}}, %{{.*}}, 0, 0
+  // CHECK: rocdl.global.load.async.to.lds.b32 %{{.*}}, %{{.*}}, 0, 0
+  // CHECK: rocdl.global.load.async.to.lds.b64 %{{.*}}, %{{.*}}, 0, 0
+  // CHECK: rocdl.global.load.async.to.lds.b128 %{{.*}}, %{{.*}}, 0, 0
+  rocdl.global.load.async.to.lds.b8 %src, %dst, 0, 0 : <1>
+  rocdl.global.load.async.to.lds.b32 %src, %dst, 0, 0 : <1>
+  rocdl.global.load.async.to.lds.b64 %src, %dst, 0, 0 : <1>
+  rocdl.global.load.async.to.lds.b128 %src, %dst, 0, 0 : <1>
+  llvm.return
+}
+
 llvm.func @rocdl.make.buffer.rsrc(%ptr : !llvm.ptr,
                                   %stride : i16,
                                   %numRecords : i64,
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 30126f6bff05a..5ae9f11360df4 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -1040,6 +1040,30 @@ llvm.func @rocdl.global.load.lds(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
   llvm.return
 }
 
+llvm.func @rocdl.global.load.async.lds.b8(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b8
+  rocdl.global.load.async.to.lds.b8 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @rocdl.global.load.async.lds.b32(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b32
+  rocdl.global.load.async.to.lds.b32 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @rocdl.global.load.async.lds.b64(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b64
+  rocdl.global.load.async.to.lds.b64 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
+llvm.func @rocdl.global.load.async.lds.b128(%src : !llvm.ptr<1>, %dst: !llvm.ptr<3>) {
+  // CHECK: call void @llvm.amdgcn.global.load.async.to.lds.b128
+  rocdl.global.load.async.to.lds.b128 %src, %dst, 0, 0 : !llvm.ptr<1>
+  llvm.return
+}
+
 llvm.func @rocdl.make.buffer.rsrc(%ptr : !llvm.ptr,
                                   %stride : i16,
                                   %numRecords : i64,

kuhar

LGTM but let's wait for @krzysz00 to confirm

lialan

LGTM.

Adds `global.load.async.to.lds` op to rocdl, supporting `b8`, `b32`, `b64` and `b128`. The op is lowered to the appropiate `llvm.amdgcn.global.load.async.to.lds.bXX` intrinsic. This is available on gfx1250+.

mlir/test/Target/LLVMIR/rocdl.mlir

ravil-mobile · 2025-10-29T10:14:15Z

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

+class ROCDL_GlobalLoadAsyncToLDSOp<string mnemonic> :
+  ROCDL_IntrOp<mnemonic, [], [], [], 0, 0, 1, 0, [2, 3], ["offset", "aux"]> {
+  dag args = (ins Arg<ROCDLGlobalBuffer, "", [MemRead]>:$globalPtr,
+                 Arg<ROCDLBufferLDS, "", [MemWrite]>:$ldsPtr,
+                 I32Attr:$offset,
+                 I32Attr:$aux);
+  let arguments = !con(args, baseArgs);
+  let assemblyFormat = [{
+    $globalPtr `,`  $ldsPtr `,` $offset `,` $aux
+    attr-dict `:` type($globalPtr)
+  }];
+  let description = [{
+    Loads data asynchronously from a global memory pointer to a local data
+    store (LDS) pointer.
+
+    Available on gfx1250+.
+  }];
+  let extraClassDefinition = [{
+    ::llvm::SmallVector<::mlir::Value> $cppClass::getAccessedOperands() {
+      return {getGlobalPtr(), getLdsPtr()};
+    }
+  }];
+}
+
+def ROCDL_GlobalLoadAsyncToLDSB8Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b8">;
+def ROCDL_GlobalLoadAsyncToLDSB32Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b32">;
+def ROCDL_GlobalLoadAsyncToLDSB64Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b64">;
+def ROCDL_GlobalLoadAsyncToLDSB128Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b128">;


NIT: Maybe we can use foreach construct in tablegen?

foreach bytes = [8, 32, 64, 128] in { let bytesStr = "b" # !cast<string>(bytes) in def ROCDL_GlobalLoadAsyncToLDS # !toupper(bytesStr) # Op : ROCDL_IntrOp<"global.load.async.to.lds." # bytesStr, [], [], [], 0, 0, 1, 0, [2, 3], ["offset", "aux"]> { dag args = (ins Arg<ROCDLGlobalBuffer, "", [MemRead]>:$globalPtr, Arg<ROCDLBufferLDS, "", [MemWrite]>:$ldsPtr, I32Attr:$offset, I32Attr:$aux); let arguments = !con(args, baseArgs); let assemblyFormat = [{ $globalPtr `,` $ldsPtr `,` $offset `,` $aux attr-dict `:` type($globalPtr) }]; let description = [{ Asynchronously loads # bytes # bytes of data from a global memory to a Local Data Store (LDS). Available on gfx1250+. }]; let extraClassDefinition = [{ ::llvm::SmallVector<::mlir::Value> $cppClass::getAccessedOperands() { return {getGlobalPtr(), getLdsPtr()}; } }]; } }

Wow, that's some cool tablegen black magic I didn't know about!

I have used foreach as you suggested, but I changed it slightly, the iterator should be bits, not bytes, and the description was not showing the right string so I also fixed that.

I have checked the generated documentation and it correctly shows the 4 ops with the right description. What I don't like is that they are not generated in order, (first goes b128, then b32, then b64, then b8).

krzysz00 · 2025-10-30T18:25:23Z

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

+  }];
+}
+
+def ROCDL_GlobalLoadAsyncToLDSB8Op : ROCDL_GlobalLoadAsyncToLDSOp<"global.load.async.to.lds.b8">;


Seconding Ravil's nits here.

pabloantoniom · 2025-11-03T16:41:54Z

Any further suggestions, or are you guys happy with the current state? @ravil-mobile @krzysz00

ravil-mobile · 2025-11-03T16:54:54Z

Any further suggestions, or are you guys happy with the current state? @ravil-mobile @krzysz00

LGTM

krzysz00

Minor note re assembly format, otherwise approved

krzysz00 · 2025-11-03T16:55:55Z

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

+    let arguments = !con(args, baseArgs);
+    let assemblyFormat = [{
+      $globalPtr `,`  $ldsPtr `,` $offset `,` $aux
+      attr-dict `:` type($globalPtr)


I'd go for type($globalPtr), type($ldsPtr) if we're doing this sort of thing

pabloantoniom requested a review from lialan October 28, 2025 11:58

pabloantoniom requested review from krzysz00 and kuhar as code owners October 28, 2025 11:58

llvmbot added mlir:llvm mlir labels Oct 28, 2025

kuhar approved these changes Oct 28, 2025

View reviewed changes

lialan approved these changes Oct 28, 2025

View reviewed changes

kuhar requested a review from ravil-mobile October 28, 2025 14:55

[mlir][rocdl] Add GlobalLoadAsyncToLDS operation

fe4f87b

Adds `global.load.async.to.lds` op to rocdl, supporting `b8`, `b32`, `b64` and `b128`. The op is lowered to the appropiate `llvm.amdgcn.global.load.async.to.lds.bXX` intrinsic. This is available on gfx1250+.

pabloantoniom force-pushed the rocdl-global-load-async branch from 0a27dd6 to fe4f87b Compare October 29, 2025 09:12

ravil-mobile reviewed Oct 29, 2025

View reviewed changes

pabloantoniom mentioned this pull request Oct 29, 2025

Add async loads support for gfx1250 ROCm/rocMLIR#2072

Open

1 task

pabloantoniom added a commit to ROCm/rocMLIR that referenced this pull request Oct 29, 2025

Cherry pick llvm/llvm-project#165374

d813c1e

pabloantoniom added a commit to ROCm/rocMLIR that referenced this pull request Oct 29, 2025

Cherry pick llvm/llvm-project#165374

2de77f3

krzysz00 reviewed Oct 30, 2025

View reviewed changes

Address ravil comments, note that I changed the foreach slightly

4a11e80

pabloantoniom requested review from krzysz00 and ravil-mobile October 31, 2025 09:15

pabloantoniom added a commit to ROCm/rocMLIR that referenced this pull request Oct 31, 2025

Cherry pick llvm/llvm-project#165374

951d3c0

krzysz00 approved these changes Nov 3, 2025

View reviewed changes

Address Krzysztofs nit

36dd422

pabloantoniom merged commit 0c65351 into llvm:main Nov 4, 2025
10 checks passed

kerbowa mentioned this pull request Nov 10, 2025

[AMDGPU] Verify dominance when rewriting spills to registers #167347

Open

[mlir][rocdl] Add GlobalLoadAsyncToLDS operation #165374

[mlir][rocdl] Add GlobalLoadAsyncToLDS operation #165374

Uh oh!

Conversation

pabloantoniom commented Oct 28, 2025

Uh oh!

llvmbot commented Oct 28, 2025

Uh oh!

llvmbot commented Oct 28, 2025

Uh oh!

kuhar left a comment

Choose a reason for hiding this comment

Uh oh!

lialan left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ravil-mobile Oct 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pabloantoniom Oct 31, 2025

Choose a reason for hiding this comment

Uh oh!

krzysz00 Oct 30, 2025

Choose a reason for hiding this comment

Uh oh!

pabloantoniom commented Nov 3, 2025

Uh oh!

ravil-mobile commented Nov 3, 2025

Uh oh!

krzysz00 left a comment

Choose a reason for hiding this comment

Uh oh!

krzysz00 Nov 3, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ravil-mobile Oct 29, 2025 •

edited

Loading