[Mlir-commits] [mlir] 760c429 - [MLIR][AMDGPU] Added l2-prefetch op to AMDGPU (#188457)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Mar 27 08:05:35 PDT 2026
Author: Ravil Dorozhinskii
Date: 2026-03-27T16:05:30+01:00
New Revision: 760c4292c041aa857a462413f07623f69276f99b
URL: https://github.com/llvm/llvm-project/commit/760c4292c041aa857a462413f07623f69276f99b
DIFF: https://github.com/llvm/llvm-project/commit/760c4292c041aa857a462413f07623f69276f99b.diff
LOG: [MLIR][AMDGPU] Added l2-prefetch op to AMDGPU (#188457)
This PR adds `global_prefetch` op to prefetch a cache line to high-level
caches using the aligned address of the source `memref` and an offset
provided by the indices of the element containing the cache line. This
provides temporal hints (e.g., regular or high-priority). Note that
out-of-bounds access is allowed in speculative mode. Ensure the source
`memref` is in address space `1`.
---------
Co-authored-by: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Added:
mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
Modified:
mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td
mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
mlir/test/Dialect/AMDGPU/invalid.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
index c862fb2fc5a3a..2f8d4a086c690 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
@@ -47,4 +47,10 @@ def AMDGPU_SchedBarrierOpOptAttr : EnumAttr<AMDGPU_Dialect, AMDGPU_SchedBarrierO
def AMDGPU_MFMAPermBAttr : EnumAttr<AMDGPU_Dialect, AMDGPU_MFMAPermB,
"mfma_perm_b">;
+def AMDGPU_TemporalHintAttr : EnumAttr<AMDGPU_Dialect, AMDGPU_TemporalHint,
+ "temporal_hint">;
+
+def AMDGPU_CacheScopeAttr : EnumAttr<AMDGPU_Dialect, AMDGPU_CacheScope,
+ "cache_scope">;
+
#endif // MLIR_DIALECT_AMDGPU_IR_AMDGPUATTRS_TD
diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td
index 4ec7cb3cd7307..89432f20575d6 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td
@@ -80,4 +80,40 @@ def AMDGPU_MFMAPermB : I32Enum<"MFMAPermB",
let cppNamespace = "::mlir::amdgpu";
}
+def AMDGPU_TemporalHint : I32Enum<"TemporalHint",
+ "AMDGPU-specific prefetch temporal hints. "
+ "RT - regular temporal for both near and far caches; "
+ "NT - non-temporal for both near and far caches; "
+ "HT - high-priority temporal for both near and far caches; "
+ "LU - last-use; "
+ "NT_RT - non-temporal for near cache(s) and regular for far caches; "
+ "RT_NT - regular for near cache(s) and non-temporal for far caches; "
+ "NT_HT - non-temporal for near cache(s) and high-priority temporal for far caches",
+ [
+ I32EnumCase<"RT", 0>,
+ I32EnumCase<"NT", 1>,
+ I32EnumCase<"HT", 2>,
+ I32EnumCase<"LU", 3>,
+ I32EnumCase<"NT_RT", 4>,
+ I32EnumCase<"RT_NT", 5>,
+ I32EnumCase<"NT_HT", 6>
+ ]> {
+ let cppNamespace = "::mlir::amdgpu";
+}
+
+def AMDGPU_CacheScope : I32Enum<"Scope",
+ "AMDGPU-specific cache scopes. "
+ "WGP - workgroup processor (CUs); "
+ "SE - shader engine (GL2); "
+ "DEV - device; "
+ "SYS - system",
+ [
+ I32EnumCase<"WGP", 0>,
+ I32EnumCase<"SE", 1>,
+ I32EnumCase<"DEV", 2>,
+ I32EnumCase<"SYS", 3>
+ ]> {
+ let cppNamespace = "::mlir::amdgpu";
+}
+
#endif // MLIR_DIALECT_AMDGPU_IR_AMDGPUENUMS_TD
diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
index 3eb039305904f..0f01a46e147f5 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td
@@ -1952,4 +1952,37 @@ def AMDGPU_DsBarrierStatePhaseParity :
}];
}
+def AMDGPU_GlobalPrefetchOp :
+ AMDGPU_Op<"global_prefetch", [MemoryEffects<[MemWrite, MemRead]>]>,
+ Arguments<(ins AnyMemRef:$src,
+ Variadic<I64>:$indices,
+ AMDGPU_TemporalHintAttr:$temporalHint,
+ AMDGPU_CacheScopeAttr:$cacheScope,
+ UnitAttr:$speculative)>,
+ Results<(outs)> {
+
+ let summary = "Prefetch data to caches.";
+ let description = [{
+ Prefetches a cache line to high-level caches using the aligned address of
+ the source `memref` and an offset provided by the indices of the element
+ containing the cache line. This provides temporal hints (e.g., regular
+ or high-priority). Note that out-of-bounds access is allowed in
+ speculative mode. The provided memref must be in the global address space
+ (`#gpu.address_space<global>` or 1).
+
+ This operation was introduced in gfx1250.
+
+ Example:
+ ```mlir
+ amdgpu.global_prefetch %src[%i, %j] RT SE speculative : memref<64x64xf16, #gpu.address_space<global>>
+ ```
+ }];
+
+ let assemblyFormat = [{
+ $src `[` $indices `]` $temporalHint $cacheScope (`speculative` $speculative^)? attr-dict `:` qualified(type($src))
+ }];
+
+ let hasVerifier = 1;
+}
+
#endif // MLIR_DIALECT_AMDGPU_IR_AMDGPUOPS_TD
diff --git a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
index 14c12f5a787a6..96d4e5da3388c 100644
--- a/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
+++ b/mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp
@@ -3950,6 +3950,56 @@ struct AMDGPUTensorLoadStoreOpLowering
}
};
+struct GlobalPrefetchOpLowering
+ : public ConvertOpToLLVMPattern<GlobalPrefetchOp> {
+ GlobalPrefetchOpLowering(const LLVMTypeConverter &converter, Chipset chipset)
+ : ConvertOpToLLVMPattern<GlobalPrefetchOp>(converter), chipset(chipset) {}
+
+ LogicalResult
+ matchAndRewrite(GlobalPrefetchOp op, GlobalPrefetchOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ if (chipset < kGfx1250)
+ return op->emitOpError("is only supported on gfx1250+");
+
+ const TemporalHint hint = op.getTemporalHint();
+ const bool isSpeculative = op.getSpeculative();
+
+ int32_t immArgValue = static_cast<int32_t>(hint);
+
+ // Note that only RT and HT can operate in both speculative and
+ // non-speculative modes. The other variants (NT_RT, RT_NT, NT_HT, etc.)
+ // operate only in the speculative mode and, therefore, do not require
+ // toggling the least significant bit for mode changes
+ // Temporal hint is encoded in lower bits - i.e. [2:0]
+ if (llvm::is_contained({TemporalHint::RT, TemporalHint::HT}, hint))
+ immArgValue = isSpeculative ? immArgValue : immArgValue | 1;
+
+ // Prefetch scope level is encoded in upper bits - i.e., [4:3]
+ immArgValue = static_cast<int32_t>(op.getCacheScope()) << 3 | immArgValue;
+
+ IntegerAttr immArgAttr = rewriter.getI32IntegerAttr(immArgValue);
+
+ ValueRange indices = adaptor.getIndices();
+ Value memRef = adaptor.getSrc();
+ MemRefDescriptor descriptor(memRef);
+ MemRefType memRefType = op.getSrc().getType();
+ Location loc = op->getLoc();
+ auto inboundsFlags = isSpeculative ? LLVM::GEPNoWrapFlags::none
+ : LLVM::GEPNoWrapFlags::inbounds |
+ LLVM::GEPNoWrapFlags::nuw;
+ Value prefetchPtr = getStridedElementPtr(
+ rewriter, loc, memRefType, descriptor, indices, inboundsFlags);
+
+ rewriter.replaceOpWithNewOp<ROCDL::GlobalPrefetchOp>(
+ op, prefetchPtr, immArgAttr, mlir::ArrayAttr{}, mlir::ArrayAttr{},
+ mlir::ArrayAttr{});
+ return success();
+ }
+
+private:
+ Chipset chipset;
+};
+
struct ConvertAMDGPUToROCDLPass
: public impl::ConvertAMDGPUToROCDLPassBase<ConvertAMDGPUToROCDLPass> {
using Base::Base;
@@ -4086,8 +4136,8 @@ void mlir::populateAMDGPUToROCDLConversionPatterns(LLVMTypeConverter &converter,
AMDGPUTensorLoadStoreOpLowering<TensorStoreFromLDSOp,
ROCDL::TensorStoreFromLDSOp>,
DsBarrierInitOpLowering, DsBarrierPollStateOpLowering,
- DsAsyncBarrierArriveOpLowering, DsBarrierArriveOpLowering>(converter,
- chipset);
+ DsAsyncBarrierArriveOpLowering, DsBarrierArriveOpLowering,
+ GlobalPrefetchOpLowering>(converter, chipset);
patterns.add<AMDGPUSwizzleBitModeLowering, DsBarrierStatePhaseOpLowering,
DsBarrierStatePendingCountOpLowering,
DsBarrierStateInitCountOpLowering,
diff --git a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
index b715f4ab93231..e27bd461908cd 100644
--- a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
+++ b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
@@ -1302,5 +1302,45 @@ LogicalResult DsBarrierArriveOp::verify() {
return verifyDsBarrierOpCommon(*this);
}
+//===----------------------------------------------------------------------===//
+// GlobalPrefetchOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GlobalPrefetchOp::verify() {
+ auto src = cast<MemRefType>(getSrc().getType());
+
+ Attribute memSpace = src.getMemorySpace();
+ if (!memSpace)
+ return this->emitOpError("the source must have address space attribute");
+ if (!hasGlobalMemorySpace(memSpace))
+ return this->emitOpError("the source must reside in global address space");
+
+ ArrayRef<int64_t> srcShape = src.getShape();
+ const size_t numIndices = getIndices().size();
+ if (srcShape.size() != numIndices)
+ return this->emitOpError(
+ "the number of indices must match the source shape size");
+
+ const TemporalHint temporalHint = getTemporalHint();
+ const bool isSpeculative = getSpeculative();
+
+ // Note that temporal hints are shared between load, store,
+ // prefetch, etc. instructions. However, some instructions
+ // operate only with a subset of hints according to the ISA
+ // documentation. In case of global prefetch, non-temporal (NT)
+ // and last-use (LU) hints are not used. The extra bits of encoding
+ // are used to encode speculative or non-speculative instruction behavior
+ if (llvm::is_contained({TemporalHint::NT, TemporalHint::LU}, temporalHint))
+ return this->emitOpError("does not support NT and LU modes");
+
+ if (llvm::is_contained(
+ {TemporalHint::NT_RT, TemporalHint::RT_NT, TemporalHint::NT_HT},
+ temporalHint) &&
+ !isSpeculative) {
+ return this->emitOpError("operates only in the speculative mode");
+ }
+ return success();
+}
+
#define GET_OP_CLASSES
#include "mlir/Dialect/AMDGPU/IR/AMDGPU.cpp.inc"
diff --git a/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir b/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
new file mode 100644
index 0000000000000..643b809d90675
--- /dev/null
+++ b/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
@@ -0,0 +1,25 @@
+// RUN: mlir-opt %s --convert-amdgpu-to-rocdl=chipset=gfx1250 --split-input-file --verify-diagnostics | FileCheck %s
+
+// CHECK-LABEL: @glb_prefetch0
+func.func @glb_prefetch0(%src : memref<64x64xf16, #gpu.address_space<global>>, %i : i64, %j : i64) {
+ // CHECK: %[[PTR:.*]] = llvm.getelementptr %{{.*}}[%{{.*}}] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, f16
+ // CHECK: rocdl.global.prefetch %{{.*}}, scope 0 : !llvm.ptr<1>
+ amdgpu.global_prefetch %src[%i, %j] RT WGP speculative : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// CHECK-LABEL: @glb_prefetch1
+func.func @glb_prefetch1(%src : memref<64x64xf16, #gpu.address_space<global>>, %i : i64, %j : i64) {
+ // CHECK: %[[PTR:.*]] = llvm.getelementptr inbounds|nuw %{{.*}}[%{{.*}}] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, f16
+ // CHECK: rocdl.global.prefetch %[[PTR]], scope 3 : !llvm.ptr<1>
+ amdgpu.global_prefetch %src[%i, %j] HT WGP : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// CHECK-LABEL: @glb_prefetch2
+func.func @glb_prefetch2(%src : memref<64x64xf16, #gpu.address_space<global>>, %i : i64, %j : i64) {
+ // CHECK: %[[PTR:.*]] = llvm.getelementptr %{{.*}}[%{{.*}}] : (!llvm.ptr<1>, i64) -> !llvm.ptr<1>, f16
+ // CHECK: rocdl.global.prefetch %[[PTR]], scope 10 : !llvm.ptr<1>
+ amdgpu.global_prefetch %src[%i, %j] HT SE speculative : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
diff --git a/mlir/test/Dialect/AMDGPU/invalid.mlir b/mlir/test/Dialect/AMDGPU/invalid.mlir
index d1bb43e5587a6..1eb64ddecf695 100644
--- a/mlir/test/Dialect/AMDGPU/invalid.mlir
+++ b/mlir/test/Dialect/AMDGPU/invalid.mlir
@@ -660,3 +660,76 @@ func.func @sparse_wmma_i4_requires_equal_length_wave64(%a: vector<8xi4>, %b: vec
%d = amdgpu.sparse_wmma 16x16x32 %a * %b + %c sparse(%idx : vector<4xi8>) {wave64} : vector<8xi4>, vector<16xi4>, vector<4xi32>
func.return %d : vector<4xi32>
}
+
+// -----
+
+// GlobalPrefetchOp: source must have address space attribute
+func.func @global_prefetch_no_address_space(%src: memref<64x64xf16>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op the source must have address space attribute}}
+ amdgpu.global_prefetch %src[%i, %j] RT WGP : memref<64x64xf16>
+ func.return
+}
+
+
+// -----
+
+// GlobalPrefetchOp: source must reside in global address space
+func.func @global_prefetch_wrong_address_space(%src: memref<64x64xf16, #gpu.address_space<workgroup>>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op the source must reside in global address space}}
+ amdgpu.global_prefetch %src[%i, %j] RT SE : memref<64x64xf16, #gpu.address_space<workgroup>>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: number of indices must match source shape rank
+func.func @global_prefetch_wrong_num_indices(%src: memref<64x64xf16, #gpu.address_space<global>>, %i: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op the number of indices must match the source shape size}}
+ amdgpu.global_prefetch %src[%i] RT DEV : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: NT temporal hint is not supported
+func.func @global_prefetch_nt_mode(%src: memref<64x64xf16, #gpu.address_space<global>>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op does not support NT and LU modes}}
+ amdgpu.global_prefetch %src[%i, %j] NT SYS : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: LU temporal hint is not supported
+func.func @global_prefetch_lu_mode(%src: memref<64x64xf16, #gpu.address_space<global>>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op does not support NT and LU modes}}
+ amdgpu.global_prefetch %src[%i, %j] LU DEV : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: NT_RT requires speculative mode
+func.func @global_prefetch_nt_rt_not_speculative(%src: memref<64x64xf16, #gpu.address_space<global>>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op operates only in the speculative mode}}
+ amdgpu.global_prefetch %src[%i, %j] NT_RT WGP : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: RT_NT requires speculative mode
+func.func @global_prefetch_rt_nt_not_speculative(%src: memref<64x64xf16, #gpu.address_space<global>>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op operates only in the speculative mode}}
+ amdgpu.global_prefetch %src[%i, %j] RT_NT SE : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: NT_HT requires speculative mode
+func.func @global_prefetch_nt_ht_not_speculative(%src: memref<64x64xf16, #gpu.address_space<global>>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op operates only in the speculative mode}}
+ amdgpu.global_prefetch %src[%i, %j] NT_HT DEV : memref<64x64xf16, #gpu.address_space<global>>
+ func.return
+}
More information about the Mlir-commits
mailing list