[ROCDL] Added global/flag data prefetch ops

[ravil-mobile]

This PR brings data prefetch ops to ROCDL for gfx1250 architecture. Extended all necessary rocdl tests

[llvmbot]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-llvm

Author: Ravil Dorozhinskii (ravil-mobile)

Changes

This PR brings data prefetch ops to ROCDL for gfx1250 architecture. Extended all necessary rocdl tests

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td (+25)
• (modified) mlir/test/Dialect/LLVMIR/rocdl.mlir (+14)
• (modified) mlir/test/Target/LLVMIR/rocdl.mlir (+14)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
index cd36300d7ac16..900cafebd95e8 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td
@@ -1173,6 +1173,31 @@ def ROCDL_RawBufferAtomicCmpSwap :
   }];
 }
 
+//===---------------------------------------------------------------------===//
+// Memory prefetch intrinsics
+
+def ROCDL_GlobalPrefetchOp :
+  ROCDL_IntrOp<"global.prefetch", [], [], [], 0, 0, 0, 0, [1], ["scope"]>,
+  Arguments<(ins Arg<LLVM_PointerInAddressSpace<1>, "", []>:$ptr, I32Attr:$scope)> {
+  let description = [{
+    Prefetches 1 byte of data per lane from global memory into the WGP-cache or L2-cache.
+    Available on gfx1250+.
+  }];
+  let results = (outs);
+  let assemblyFormat = "$ptr `scope` $scope attr-dict";
+}
+
+def ROCDL_FlatPrefetchOp :
+  ROCDL_IntrOp<"flat.prefetch", [], [], [], 0, 0, 0, 0, [1], ["scope"]>,
+  Arguments<(ins Arg<LLVM_PointerInAddressSpace<0>, "", []>:$ptr, I32Attr:$scope)> {
+  let description = [{
+    Prefetches 1 byte of data per lane using flat-memory addresses into the WGP-cache or L2-cache.
+    Available on gfx1250+.
+  }];
+  let results = (outs);
+  let assemblyFormat = "$ptr `scope` $scope attr-dict";
+}
+
 //===---------------------------------------------------------------------===//
 // MI-100 and MI-200 buffer atomic floating point add intrinsic
 
diff --git a/mlir/test/Dialect/LLVMIR/rocdl.mlir b/mlir/test/Dialect/LLVMIR/rocdl.mlir
index 40084bc07d4f7..3e048b93307fc 100644
--- a/mlir/test/Dialect/LLVMIR/rocdl.mlir
+++ b/mlir/test/Dialect/LLVMIR/rocdl.mlir
@@ -878,6 +878,20 @@ llvm.func @rocdl.raw.ptr.buffer.i32(%rsrc : !llvm.ptr<8>,
   llvm.return
 }
 
+llvm.func @rocdl.global.prefetch(%ptr : !llvm.ptr<1>) {
+  // CHECK-LABEL: rocdl.global.prefetch
+  // CHECK: rocdl.global.prefetch %{{.*}} scope 0
+  rocdl.global.prefetch %ptr scope 0
+  llvm.return
+}
+
+llvm.func @rocdl.flat.prefetch(%ptr : !llvm.ptr<0>) {
+  // CHECK-LABEL: rocdl.flat.prefetch
+  // CHECK: rocdl.flat.prefetch %{{.*}} scope 0
+  rocdl.flat.prefetch %ptr scope 0
+  llvm.return
+}
+
 // -----
 
 llvm.func @rocdl.raw.buffer.f32(%rsrc : vector<4xi32>,
diff --git a/mlir/test/Target/LLVMIR/rocdl.mlir b/mlir/test/Target/LLVMIR/rocdl.mlir
index 2c748ad509356..5e8ce60415073 100644
--- a/mlir/test/Target/LLVMIR/rocdl.mlir
+++ b/mlir/test/Target/LLVMIR/rocdl.mlir
@@ -1341,6 +1341,20 @@ llvm.func @rocdl.raw.ptr.buffer.load.lds(%rsrc : !llvm.ptr<8>, %dstLds : !llvm.p
   llvm.return
 }
 
+llvm.func @rocdl.global.prefetch(%ptr : !llvm.ptr<1>) {
+  // CHECK-LABEL: rocdl.global.prefetch
+  // CHECK: call void @llvm.amdgcn.global.prefetch(ptr addrspace(1) %{{.*}}, i32 0)
+  rocdl.global.prefetch %ptr scope 0
+  llvm.return
+}
+
+llvm.func @rocdl.flat.prefetch(%ptr : !llvm.ptr<0>) {
+  // CHECK-LABEL: rocdl.flat.prefetch
+  // CHECK: call void @llvm.amdgcn.flat.prefetch(ptr %{{.*}}, i32 0)
+  rocdl.flat.prefetch %ptr scope 0
+  llvm.return
+}
+
 llvm.func @rocdl.wmma.scale(%arg0: i32, %arg1: vector<4xf32>, %arg2: vector<8xi32>,
                             %arg3: vector<12xi32>, %arg5: vector<16xi32>,
                             %arg8: i64, %arg9: vector<8xf32>) -> vector<4xf32> {

[kuhar]

Should we also print the pointer type?

[ravil-mobile]

We usually provide types for the results. These ops do not have return type. @krzysz00, what do you think?

[kuhar]

IMO types help readability, especially in low-level IR. We often omit operand types if they match result types, but here there's no result, so I'd prefer to have it over the operand instead.

[krzysz00]

Yeah, I think I'm with @kuhar on this one these days. Not enough to actively go on a syntax change rampage, but we probably should put the types in here

[ravil-mobile]

Ok, added types there

[kuhar]

also here

[krzysz00]

Once Jakub's comments are addressed, LGTM

[kuhar]

Suggested change

	let assemblyFormat = "$ptr qualified(type($ptr)) `,` `scope` $scope attr-dict";
	let assemblyFormat = "$ptr `,` `scope` $scope attr-dict : qualified(type($ptr))";

?

[ravil-mobile]

Reworked. I didn't want to do it at the first place because I wanted that it was clear that there was not return type. But anyway, your proposal is some standard standard. So, in short, I agree and reworked.

[kuhar]

Thanks for the changes