[Mlir-commits] [mlir] [MLIR][AMDGPU] Added l2-prefetch op to AMDGPU (PR #188457)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Mar 25 03:37:58 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
@llvm/pr-subscribers-mlir-amdgpu
Author: Ravil Dorozhinskii (ravil-mobile)
<details>
<summary>Changes</summary>
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`.
---
Full diff: https://github.com/llvm/llvm-project/pull/188457.diff
7 Files Affected:
- (modified) mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td (+3)
- (modified) mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td (+21)
- (modified) mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUOps.td (+33)
- (modified) mlir/lib/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.cpp (+52-2)
- (modified) mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp (+30)
- (added) mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir (+15)
- (modified) mlir/test/Dialect/AMDGPU/invalid.mlir (+54)
``````````diff
diff --git a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
index c862fb2fc5a3a..0e4ab8d5b6dc5 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUAttrs.td
@@ -47,4 +47,7 @@ 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">;
+
#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..68bae3d255447 100644
--- a/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td
+++ b/mlir/include/mlir/Dialect/AMDGPU/IR/AMDGPUEnums.td
@@ -80,4 +80,25 @@ 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; ",
+ [
+ I32EnumAttrCase<"RT", 0>,
+ I32EnumAttrCase<"NT", 1>,
+ I32EnumAttrCase<"HT", 2>,
+ I32EnumAttrCase<"LU", 3>,
+ I32EnumAttrCase<"NT_RT", 4>,
+ I32EnumAttrCase<"RT_NT", 5>,
+ I32EnumAttrCase<"NT_HT", 6>
+ ]> {
+ 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..22d0ec9822ea2 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,
+ 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. Ensure the source memref is in address space `1`.
+
+ This operation was introduced in gfx1250.
+
+ Example:
+ ```mlir
+ amdgpu.global_prefetch %src[%i, %j] RT speculative : memref<64x64xf16, 1>
+ ```
+ }];
+
+ let assemblyFormat = [{
+ $src `[` $indices `]` $temporalHint (`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..9667a081d1ea1 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 llvmScopeValue = static_cast<int32_t>(hint);
+ if ((hint == TemporalHint::RT) || (hint == TemporalHint::HT))
+ llvmScopeValue = isSpeculative ? llvmScopeValue : llvmScopeValue | 1;
+
+ IntegerAttr scopeAttr = rewriter.getI32IntegerAttr(llvmScopeValue);
+
+ ValueRange indices = adaptor.getIndices();
+ Value memRef = adaptor.getSrc();
+ MemRefDescriptor descriptor(memRef);
+ Location loc = op->getLoc();
+ Value offset =
+ LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(), 0);
+ for (size_t i = 0; i < indices.size(); ++i) {
+ Value stride = descriptor.stride(rewriter, loc, i);
+ Value mulOp = LLVM::MulOp::create(rewriter, loc, rewriter.getI64Type(),
+ stride, indices[i]);
+ offset = LLVM::AddOp::create(rewriter, loc, rewriter.getI64Type(), offset,
+ mulOp);
+ }
+
+ Value basePtr = descriptor.alignedPtr(rewriter, loc);
+ Type elemTy = op.getSrc().getType().getElementType();
+ Type llvmElemTy = getTypeConverter()->convertType(elemTy);
+ Value prefetchPtr = LLVM::GEPOp::create(rewriter, loc, basePtr.getType(),
+ llvmElemTy, basePtr, offset);
+ Operation *newOp = ROCDL::GlobalPrefetchOp::create(
+ rewriter, loc, prefetchPtr, scopeAttr, {}, {}, {});
+
+ rewriter.replaceOp(op, newOp);
+ 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..ebbdbf0b07a55 100644
--- a/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
+++ b/mlir/lib/Dialect/AMDGPU/IR/AMDGPUOps.cpp
@@ -1302,5 +1302,35 @@ LogicalResult DsBarrierArriveOp::verify() {
return verifyDsBarrierOpCommon(*this);
}
+//===----------------------------------------------------------------------===//
+// GlobalPrefetchOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult GlobalPrefetchOp::verify() {
+ auto src = cast<MemRefType>(getSrc().getType());
+
+ const unsigned memorySpace = src.getMemorySpaceAsInt();
+ if (memorySpace != 1)
+ return this->emitOpError("the source must reside in address space `1`");
+
+ 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();
+ if (temporalHint == TemporalHint::NT)
+ return this->emitOpError("does not support NT mode");
+ if ((temporalHint == TemporalHint::NT_RT) ||
+ (temporalHint == TemporalHint::RT_NT) ||
+ (temporalHint == TemporalHint::NT_HT)) {
+ if (!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..f8e9db5730b6b
--- /dev/null
+++ b/mlir/test/Conversion/AMDGPUToROCDL/global-prefetch.mlir
@@ -0,0 +1,15 @@
+// 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, 1>, %i : i64, %j : i64) {
+ // CHECK: rocdl.global.prefetch %{{.*}}, scope 0 : !llvm.ptr<1>
+ amdgpu.global_prefetch %src[%i, %j] RT speculative : memref<64x64xf16, 1>
+ func.return
+}
+
+// CHECK-LABEL: @glb_prefetch1
+func.func @glb_prefetch1(%src : memref<64x64xf16, 1>, %i : i64, %j : i64) {
+ // CHECK: rocdl.global.prefetch %{{.*}}, scope 3 : !llvm.ptr<1>
+ amdgpu.global_prefetch %src[%i, %j] HT : memref<64x64xf16, 1>
+ func.return
+}
diff --git a/mlir/test/Dialect/AMDGPU/invalid.mlir b/mlir/test/Dialect/AMDGPU/invalid.mlir
index d1bb43e5587a6..595c00c2c1f4c 100644
--- a/mlir/test/Dialect/AMDGPU/invalid.mlir
+++ b/mlir/test/Dialect/AMDGPU/invalid.mlir
@@ -660,3 +660,57 @@ 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 reside in address space 1
+func.func @global_prefetch_wrong_address_space(%src: memref<64x64xf16>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op the source must reside in address space `1`}}
+ amdgpu.global_prefetch %src[%i, %j] RT : memref<64x64xf16>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: number of indices must match source shape rank
+func.func @global_prefetch_wrong_num_indices(%src: memref<64x64xf16, 1>, %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 : memref<64x64xf16, 1>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: NT temporal hint is not supported
+func.func @global_prefetch_nt_mode(%src: memref<64x64xf16, 1>, %i: i64, %j: i64) {
+ // expected-error at +1 {{'amdgpu.global_prefetch' op does not support NT mode}}
+ amdgpu.global_prefetch %src[%i, %j] NT : memref<64x64xf16, 1>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: NT_RT requires speculative mode
+func.func @global_prefetch_nt_rt_not_speculative(%src: memref<64x64xf16, 1>, %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 : memref<64x64xf16, 1>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: RT_NT requires speculative mode
+func.func @global_prefetch_rt_nt_not_speculative(%src: memref<64x64xf16, 1>, %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 : memref<64x64xf16, 1>
+ func.return
+}
+
+// -----
+
+// GlobalPrefetchOp: NT_HT requires speculative mode
+func.func @global_prefetch_nt_ht_not_speculative(%src: memref<64x64xf16, 1>, %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 : memref<64x64xf16, 1>
+ func.return
+}
``````````
</details>
https://github.com/llvm/llvm-project/pull/188457
More information about the Mlir-commits
mailing list