[Mlir-commits] [mlir] [xegpu][transform] Improve `insert_prefetch` to support loads from dynamic shaped memrefs. (PR #228258)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 1 14:49:41 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Charitha Saumya (charithaintc)
<details>
<summary>Changes</summary>
Current `xegpu.insert_prefetch` implicitly assumes that the provided load op's descriptor is outside the `scf.for`. This is not true for dynamic shapes. Dynamic shapes usually require creating a subview inside the loop and creating a new descriptor off of the subview. This PR attempts to support this case.
Alternative design:
It may also be possible to remove this `subview` by folding it into the load op's offsets. This approach is not attempted and maybe more cleaner because it gives a nice canonical form without subviews. I welcome any discussion on this design choice.
Assisted By: Codex
---
Patch is 22.22 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/228258.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td (+3-2)
- (modified) mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp (+175-20)
- (modified) mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir (+139)
- (modified) mlir/test/Dialect/XeGPU/transform-ops.mlir (+56)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 2ed4690129708..5ef0c25e003d3 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -152,8 +152,9 @@ def InsertPrefetchOp : Op<Transform_Dialect, "xegpu.insert_prefetch", [
let description = [{
Inserts `xegpu.prefetch_nd` operations for the given `xegpu.load_nd` op.
The load op must reside within the `scf.for` loop. Number of prefetch steps
- is set by the `nb_prefetch` argument (default value is 1). Returns a handle
- to the created `xegpu.create_nd_desc` op.
+ is set by the `nb_prefetch` argument (default value is 1).
+
+ Returns a handle to the created `xegpu.create_nd_tdesc` ops.
}];
let arguments = (ins TransformHandleTypeInterface:$target,
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index d10265772cacb..7ecc3e5e42e0c 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -8,10 +8,14 @@
#include "mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include <optional>
@@ -348,6 +352,48 @@ void transform::SetGPULaunchThreadsOp::getEffects(
modifiesPayload(effects);
}
+/// Collect the pure loop-local producer slice (i.e. set of ops) needed to
+/// recreate a descriptor for another iteration. Values defined above the loop
+/// are reused.
+static LogicalResult
+collectLoopLocalDescriptorDependancyChain(Value value, scf::ForOp forOp,
+ SmallPtrSetImpl<Operation *> &seen,
+ SmallVectorImpl<Operation *> &slice) {
+ if (auto arg = dyn_cast<BlockArgument>(value))
+ return forOp.isDefinedOutsideOfLoop(arg) || arg == forOp.getInductionVar()
+ ? success()
+ : failure();
+
+ Operation *def = value.getDefiningOp();
+ if (forOp.isDefinedOutsideOfLoop(value))
+ return success();
+ if (def->getBlock() != forOp.getBody() || def->getNumRegions() != 0 ||
+ !isMemoryEffectFree(def))
+ return failure();
+ if (!seen.insert(def).second)
+ return success();
+ for (Value operand : def->getOperands())
+ if (failed(collectLoopLocalDescriptorDependancyChain(operand, forOp, seen,
+ slice)))
+ return failure();
+ slice.push_back(def);
+ return success();
+}
+
+/// Clone the loop-local producer slice with the current induction variable
+/// replaced by the iteration that should be prefetched.
+static xegpu::CreateNdDescOp
+clonePrefetchDescriptor(transform::TransformRewriter &rewriter,
+ scf::ForOp forOp, xegpu::CreateNdDescOp descOp,
+ ArrayRef<Operation *> slice, Value iteration) {
+ IRMapping mapping;
+ mapping.map(forOp.getInductionVar(), iteration);
+ for (Operation *op : slice)
+ rewriter.clone(*op, mapping);
+ return mapping.lookup(descOp.getResult())
+ .getDefiningOp<xegpu::CreateNdDescOp>();
+}
+
DiagnosedSilenceableFailure
transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
transform::TransformResults &results,
@@ -399,26 +445,109 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
return diag;
}
- // Find descriptor op.
+ // Find the load's descriptor and distinguish the two supported forms:
+ // (1) a descriptor outside the loop with the loop induction variable in a
+ // load offset, or
+ // (2) a descriptor is created off of a subview. This subview must carry the
+ // loop induction variable in its offsets. It is expected that load
+ // offsets are zero.
auto maybeDescOp =
findProducerOfType<xegpu::CreateNdDescOp>(loadOp.getResult());
if (!maybeDescOp)
return emitSilenceableFailure(getLoc()) << "Could not find descriptor op.";
auto descOp = *maybeDescOp;
+ bool loopLocalDesc = forOp->isAncestor(descOp);
+ SmallVector<OpFoldResult> loadOffsets = loadOp.getMixedOffsets();
+ if (loopLocalDesc) {
+ // Case 2: the loop position must come from a subview, not load_nd.
+ if (!areAllConstantIntValue(loadOffsets, 0))
+ return emitSilenceableFailure(getLoc())
+ << "loop-local prefetch requires zero load offsets";
+ } else {
+ // Case 1: at least one load offset must be the induction variable itself.
+ // Every other offset must already be available before the loop.
+ bool loadUsesIV = false;
+ for (OpFoldResult offset : loadOffsets) {
+ auto value = dyn_cast<Value>(offset);
+ if (!value)
+ continue;
+ if (value == forOp.getInductionVar()) {
+ loadUsesIV = true;
+ continue;
+ }
+ if (!forOp.isDefinedOutsideOfLoop(value))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch load offsets must use the induction variable "
+ "directly or be loop-invariant";
+ }
+ if (!loadUsesIV)
+ return emitSilenceableFailure(getLoc())
+ << "prefetch load offsets must directly use the loop induction "
+ "variable";
+ }
+
+ SmallVector<Operation *> descSlice;
+ if (loopLocalDesc) {
+ // Case 2: cloning a future subview needs a boundary-checking descriptor
+ // and a side-effect-free chain of producers from within the loop body.
+ if (!descOp.getType().getBoundaryCheck())
+ return emitSilenceableFailure(getLoc())
+ << "loop-local prefetch requires descriptor boundary_check";
- // Clone desc op outside the loop.
+ SmallPtrSet<Operation *, 8> seen;
+ if (failed(collectLoopLocalDescriptorDependancyChain(
+ descOp.getResult(), forOp, seen, descSlice)))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch descriptor depends on unsupported loop-local values";
+
+ // A local subview must carry the induction variable directly in an
+ // offset. Its remaining offsets must be reusable in the prologue.
+ bool hasSubview = false;
+ bool subviewUsesIV = false;
+ for (Operation *op : descSlice) {
+ auto subview = dyn_cast<memref::SubViewOp>(op);
+ if (!subview)
+ continue;
+ hasSubview = true;
+ for (OpFoldResult offset : subview.getMixedOffsets()) {
+ auto value = dyn_cast<Value>(offset);
+ if (!value)
+ continue;
+ if (value == forOp.getInductionVar()) {
+ subviewUsesIV = true;
+ continue;
+ }
+ if (!forOp.isDefinedOutsideOfLoop(value))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch subview offsets must use the induction variable "
+ "directly or be loop-invariant";
+ }
+ }
+ if (!hasSubview || !subviewUsesIV)
+ return emitSilenceableFailure(getLoc())
+ << "loop-local prefetch requires a subview offset that directly "
+ "uses the loop induction variable";
+ }
+
+ // Case 1 shares one descriptor across all prefetches. Case 2 recreates its
+ // loop-local descriptor at each prefetch site below.
rewriter.setInsertionPoint(forOp);
- auto newDescOp =
- cast<xegpu::CreateNdDescOp>(rewriter.clone(*descOp.getOperation()));
+ xegpu::CreateNdDescOp newDescOp;
+ if (!loopLocalDesc)
+ newDescOp =
+ cast<xegpu::CreateNdDescOp>(rewriter.clone(*descOp.getOperation()));
- // Clone reduction loop to emit initial prefetches.
- // Compute upper bound of the init loop: start + nbPrefetch * step.
+ // Warm up the first nbPrefetch tiles before the load loop. Case 2 clamps
+ // this prologue to the runtime bound so its subviews remain in-bounds.
auto nbPrefetchCst =
arith::ConstantIndexOp::create(rewriter, forOp.getLoc(), nbPrefetch);
auto nbStep = rewriter.createOrFold<arith::MulIOp>(
forOp.getLoc(), nbPrefetchCst, forOp.getStep());
auto initUpBound = rewriter.createOrFold<arith::AddIOp>(
forOp.getLoc(), forOp.getLowerBound(), nbStep);
+ if (loopLocalDesc)
+ initUpBound = arith::MinSIOp::create(rewriter, forOp.getLoc(), initUpBound,
+ forOp.getUpperBound());
auto initForOp =
scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(),
initUpBound, forOp.getStep());
@@ -427,10 +556,13 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
auto readCacheHint =
xegpu::CachePolicyAttr::get(ctx, xegpu::CachePolicy::CACHED);
- // Modify loadOp mixedOffsets by replacing the for loop induction variable
- // with the given value.
+ // Case 1 substitutes the iteration into load_nd's direct-IV offset. In
+ // case 2 the cloned subview carries the iteration, so load offsets stay zero.
auto getPrefetchOffsets =
[&](Value replacementVal) -> SmallVector<OpFoldResult> {
+ if (loopLocalDesc)
+ return getAsIndexOpFoldResult(
+ ctx, SmallVector<int64_t>(loadOffsets.size(), 0));
IRMapping mapping;
mapping.map(forOp.getInductionVar(), replacementVal);
SmallVector<Value> dynamicOffsets =
@@ -441,31 +573,54 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
return getMixedValues(constOffsets, dynamicOffsets, ctx);
};
- // Insert prefetch op in init loop.
- // Replace induction var with the init loop induction var.
+ // Issue warmup prefetches. Case 2 clones the subview and descriptor with the
+ // prologue induction variable replacing the load loop's variable.
rewriter.setInsertionPointToStart(initForOp.getBody());
- xegpu::PrefetchNdOp::create(rewriter, newDescOp.getLoc(),
- newDescOp.getResult(),
+ xegpu::CreateNdDescOp initDescOp =
+ loopLocalDesc
+ ? clonePrefetchDescriptor(rewriter, forOp, descOp, descSlice,
+ initForOp.getInductionVar())
+ : newDescOp;
+ xegpu::PrefetchNdOp::create(rewriter, initDescOp.getLoc(),
+ initDescOp.getResult(),
getPrefetchOffsets(initForOp.getInductionVar()),
readCacheHint, readCacheHint, readCacheHint,
/*layout=*/nullptr);
- // Insert prefetch op in main loop.
- // Calculate prefetch offset after the init prefetches have been issued.
+ // Prefetch nbPrefetch steps ahead at the start of each load-loop iteration.
rewriter.setInsertionPointToStart(forOp.getBody());
auto prefetchOffset = arith::AddIOp::create(rewriter, forOp.getLoc(),
forOp.getInductionVar(), nbStep);
- // Replace induction var with correct offset.
- xegpu::PrefetchNdOp::create(rewriter, newDescOp.getLoc(),
- newDescOp.getResult(),
+ xegpu::CreateNdDescOp mainDescOp = newDescOp;
+ if (loopLocalDesc) {
+ // A descriptor's boundary check does not make an out-of-bounds subview
+ // valid. Only build the future subview while its iteration is in range.
+ auto inBounds = arith::CmpIOp::create(
+ rewriter, forOp.getLoc(), arith::CmpIPredicate::slt, prefetchOffset,
+ forOp.getUpperBound());
+ auto ifOp = scf::IfOp::create(rewriter, forOp.getLoc(), inBounds,
+ /*withElseRegion=*/false);
+ rewriter.setInsertionPointToStart(ifOp.thenBlock());
+ mainDescOp = clonePrefetchDescriptor(rewriter, forOp, descOp, descSlice,
+ prefetchOffset);
+ }
+ // Replace induction var with the future tile's offset.
+ xegpu::PrefetchNdOp::create(rewriter, mainDescOp.getLoc(),
+ mainDescOp.getResult(),
getPrefetchOffsets(prefetchOffset), readCacheHint,
readCacheHint, readCacheHint, /*layout=*/nullptr);
- // Unroll the init loop.
- if (failed(loopUnrollFull(initForOp)))
+ // Case 1 has a fixed-count prologue that can be unrolled. Case 2 keeps its
+ // runtime-bounded prologue as a loop.
+ if (!loopLocalDesc && failed(loopUnrollFull(initForOp)))
return emitSilenceableFailure(getLoc()) << "Failed to unroll the loop";
- results.set(llvm::cast<OpResult>(getResult()), {newDescOp});
+ // Return both descriptors for case 2 so callers can annotate both prefetch
+ // sites; case 1 has only the shared descriptor.
+ if (loopLocalDesc)
+ results.set(llvm::cast<OpResult>(getResult()), {initDescOp, mainDescOp});
+ else
+ results.set(llvm::cast<OpResult>(getResult()), {newDescOp});
return DiagnosedSilenceableFailure::success();
}
diff --git a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
index ba259f311d76e..886455ff5995b 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
@@ -101,3 +101,142 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+func.func @insert_prefetch_loop_local_offset(%arg0: memref<128x?xf16>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %k = memref.dim %arg0, %c1 : memref<128x?xf16>
+ scf.for %iv = %c0 to %k step %c16 {
+ %remaining = arith.subi %k, %iv : index
+ %size = arith.minsi %remaining, %c16 : index
+ %tile = memref.subview %arg0[0, %iv] [128, %size] [1, 1]
+ : memref<128x?xf16> to memref<128x?xf16, strided<[?, 1], offset: ?>>
+ %desc = xegpu.create_nd_tdesc %tile
+ : memref<128x?xf16, strided<[?, 1], offset: ?>>
+ -> !xegpu.tensor_desc<128x16xf16>
+ %offset = arith.addi %iv, %c0 : index
+ %load = xegpu.load_nd %desc[0, %offset]
+ : !xegpu.tensor_desc<128x16xf16> -> vector<128x16xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{loop-local prefetch requires zero load offsets}}
+ %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @insert_prefetch_no_iv(%arg0: memref<128x128xf16>) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c128 = arith.constant 128 : index
+ %desc = xegpu.create_nd_tdesc %arg0
+ : memref<128x128xf16> -> !xegpu.tensor_desc<16x16xf16>
+ scf.for %iv = %c0 to %c128 step %c16 {
+ %load = xegpu.load_nd %desc[0, 0]
+ : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{prefetch load offsets must directly use the loop induction variable}}
+ %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @insert_prefetch_derived_load_offset(%arg0: memref<128x128xf16>) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c128 = arith.constant 128 : index
+ %desc = xegpu.create_nd_tdesc %arg0
+ : memref<128x128xf16> -> !xegpu.tensor_desc<16x16xf16>
+ scf.for %iv = %c0 to %c128 step %c16 {
+ %offset = arith.addi %iv, %c0 : index
+ %load = xegpu.load_nd %desc[%offset, %iv]
+ : !xegpu.tensor_desc<16x16xf16> -> vector<16x16xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{prefetch load offsets must use the induction variable directly or be loop-invariant}}
+ %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @insert_prefetch_no_subview_iv(%arg0: memref<128x?xf16>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %k = memref.dim %arg0, %c1 : memref<128x?xf16>
+ scf.for %iv = %c0 to %k step %c16 {
+ %tile = memref.subview %arg0[0, 0] [128, 16] [1, 1]
+ : memref<128x?xf16> to memref<128x16xf16, strided<[?, 1]>>
+ %desc = xegpu.create_nd_tdesc %tile
+ : memref<128x16xf16, strided<[?, 1]>>
+ -> !xegpu.tensor_desc<128x16xf16>
+ %load = xegpu.load_nd %desc[0, 0]
+ : !xegpu.tensor_desc<128x16xf16> -> vector<128x16xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{loop-local prefetch requires a subview offset that directly uses the loop induction variable}}
+ %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @insert_prefetch_derived_subview_offset(%arg0: memref<128x?xf16>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %k = memref.dim %arg0, %c1 : memref<128x?xf16>
+ scf.for %iv = %c0 to %k step %c16 {
+ %offset = arith.addi %iv, %c0 : index
+ %remaining = arith.subi %k, %iv : index
+ %size = arith.minsi %remaining, %c16 : index
+ %tile = memref.subview %arg0[0, %offset] [128, %size] [1, 1]
+ : memref<128x?xf16> to memref<128x?xf16, strided<[?, 1], offset: ?>>
+ %desc = xegpu.create_nd_tdesc %tile
+ : memref<128x?xf16, strided<[?, 1], offset: ?>>
+ -> !xegpu.tensor_desc<128x16xf16>
+ %load = xegpu.load_nd %desc[0, 0]
+ : !xegpu.tensor_desc<128x16xf16> -> vector<128x16xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %load = transform.structured.match ops{["xegpu.load_nd"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ // expected-error at below {{prefetch subview offsets must use the induction variable directly or be loop-invariant}}
+ %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 1 : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/transform-ops.mlir b/mlir/test/Dialect/XeGPU/transform-ops.mlir
index c6d4f8cd28c65..c0b00f72d28e2 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -433,6 +433,62 @@ module attributes {transform.with_named_sequence} {
// -----
+// CHECK-LABEL: @insert_prefetch_dynamic
+func.func @insert_prefetch_dynamic(%arg0: memref<128x?xf16>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ // CHECK: %[[K:.*]] = memref.dim %arg0
+ %k = memref.dim %arg0, %c1 : memref<128x?xf16>
+ // CHECK: %[[OUTER:.*]] = memref.subview %arg0[0, 0]
+ %outer = memref.subview %arg0[0, 0] [128, %k] [1, 1]
+ : memref<128x?xf16> to memref<128x?xf16, strided<[?, 1]>>
+ // CHECK: %[[INIT_END:.*]] = arith.minsi {{.*}}, %[[K]]
+ // CHECK: scf.for %[[INIT_IV:.*]] = %{{.*}} to %[[INIT_E...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/228258
More information about the Mlir-commits
mailing list