[Mlir-commits] [mlir] [xegpu][transform] Improve `insert_prefetch` to support loads from dynamic shaped memrefs. (PR #228258)
Charitha Saumya
llvmlistbot at llvm.org
Thu Oct 1 14:49:02 PDT 2026
https://github.com/charithaintc created https://github.com/llvm/llvm-project/pull/228258
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.
>From 1a92c04899c535c846e8f9e2a41963550f9afb9f Mon Sep 17 00:00:00 2001
From: Charitha Saumya <charitha.saumya.gusthinna.waduge at intel.com>
Date: Thu, 1 Oct 2026 20:01:04 +0000
Subject: [PATCH 1/2] add initial support for dynamic subviews
---
.../XeGPU/TransformOps/XeGPUTransformOps.td | 12 ++-
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 91 ++++++++++++++++++-
.../Dialect/XeGPU/transform-ops-invalid.mlir | 31 +++++++
mlir/test/Dialect/XeGPU/transform-ops.mlir | 43 +++++++++
4 files changed, 172 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 2ed4690129708..1559648c356e3 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -152,8 +152,16 @@ 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).
+
+ For a descriptor made from loop-local, rank-preserving, unit-stride
+ `memref.subview` ops, the prefetch uses a descriptor of the full source
+ memref instead. Subview offsets are added to the prefetch offsets. This
+ allows future offsets past the source bounds, as in the static case; the
+ descriptor must have `boundary_check = true`. Loop-local offsets may only
+ use the loop induction variable directly.
+
+ Returns a handle to the created `xegpu.create_nd_tdesc` op.
}];
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..5771dc1a0292d 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -7,7 +7,9 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.h"
+#include "mlir/Dialect/Arith/Utils/Utils.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/XeGPU/IR/XeGPU.h"
@@ -405,11 +407,77 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
if (!maybeDescOp)
return emitSilenceableFailure(getLoc()) << "Could not find descriptor op.";
auto descOp = *maybeDescOp;
+ bool loopLocalDesc = forOp->isAncestor(descOp);
+ Value prefetchSource = descOp.getSource();
+ SmallVector<SmallVector<OpFoldResult>> subviewOffsets;
+ if (loopLocalDesc) {
+ if (!descOp.getType().getBoundaryCheck())
+ return emitSilenceableFailure(getLoc())
+ << "loop-local prefetch requires descriptor boundary_check";
+
+ // Drop rank-preserving, unit-stride subviews made inside the loop. Their
+ // offsets move to prefetch_nd, while the descriptor uses the full source
+ // memref and remains valid for an out-of-bounds future tile.
+ while (Operation *def = prefetchSource.getDefiningOp()) {
+ if (!forOp->isAncestor(def))
+ break;
+ auto subview = dyn_cast<memref::SubViewOp>(def);
+ if (!subview ||
+ subview.getSourceType().getRank() != subview.getType().getRank() ||
+ !areAllConstantIntValue(subview.getMixedStrides(), 1))
+ return emitSilenceableFailure(getLoc())
+ << "loop-local prefetch requires rank-preserving unit-stride "
+ "subviews";
+ subviewOffsets.push_back(subview.getMixedOffsets());
+ prefetchSource = subview.getSource();
+ }
+ if (auto arg = dyn_cast<BlockArgument>(prefetchSource);
+ arg && arg.getOwner() == forOp.getBody())
+ return emitSilenceableFailure(getLoc())
+ << "prefetch source depends on a loop iter_arg";
+
+ // Only the induction variable itself can be replaced at both the
+ // prologue and main-loop prefetch sites.
+ auto supportedOffset = [&](Value offset) {
+ if (offset == forOp.getInductionVar())
+ return true;
+ if (Operation *def = offset.getDefiningOp();
+ def && forOp->isAncestor(def))
+ return false;
+ if (auto arg = dyn_cast<BlockArgument>(offset);
+ arg && arg.getOwner() == forOp.getBody())
+ return false;
+ return true;
+ };
+ for (Value offset : loadOp.getOffsets())
+ if (!supportedOffset(offset))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch load offset depends on unsupported loop-local "
+ "values";
+ for (const auto &offsets : subviewOffsets)
+ for (OpFoldResult offset : offsets)
+ if (auto value = dyn_cast<Value>(offset);
+ value && !supportedOffset(value))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch subview offset depends on unsupported loop-local "
+ "values";
+ if (!isa<MemRefType>(prefetchSource.getType()))
+ return emitSilenceableFailure(getLoc())
+ << "loop-local prefetch requires a memref source";
+ if (auto *def = prefetchSource.getDefiningOp();
+ def && forOp->isAncestor(def))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch source depends on unsupported loop-local values";
+ }
- // Clone desc op outside the loop.
+ // The prefetch descriptor uses a source that dominates the loop.
rewriter.setInsertionPoint(forOp);
auto newDescOp =
- cast<xegpu::CreateNdDescOp>(rewriter.clone(*descOp.getOperation()));
+ loopLocalDesc
+ ? xegpu::CreateNdDescOp::create(
+ rewriter, descOp.getLoc(), descOp.getType(),
+ cast<TypedValue<MemRefType>>(prefetchSource))
+ : cast<xegpu::CreateNdDescOp>(rewriter.clone(*descOp.getOperation()));
// Clone reduction loop to emit initial prefetches.
// Compute upper bound of the init loop: start + nbPrefetch * step.
@@ -433,6 +501,23 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
[&](Value replacementVal) -> SmallVector<OpFoldResult> {
IRMapping mapping;
mapping.map(forOp.getInductionVar(), replacementVal);
+ if (loopLocalDesc) {
+ auto mappedValue = [&](OpFoldResult offset) {
+ if (auto value = dyn_cast<Value>(offset))
+ return mapping.lookupOrDefault(value);
+ return getValueOrCreateConstantIndexOp(rewriter, forOp.getLoc(),
+ offset);
+ };
+ SmallVector<OpFoldResult> offsets;
+ for (auto [d, loadOffset] : llvm::enumerate(loadOp.getMixedOffsets())) {
+ Value combined = mappedValue(loadOffset);
+ for (const auto &subview : subviewOffsets)
+ combined = rewriter.createOrFold<arith::AddIOp>(
+ forOp.getLoc(), combined, mappedValue(subview[d]));
+ offsets.push_back(combined);
+ }
+ return offsets;
+ }
SmallVector<Value> dynamicOffsets =
llvm::map_to_vector(loadOp.getOffsets(), [&](Value v) {
return mapping.lookupOrDefault(v);
@@ -455,7 +540,7 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
rewriter.setInsertionPointToStart(forOp.getBody());
auto prefetchOffset = arith::AddIOp::create(rewriter, forOp.getLoc(),
forOp.getInductionVar(), nbStep);
- // Replace induction var with correct offset.
+ // Replace induction var with the future tile's offset.
xegpu::PrefetchNdOp::create(rewriter, newDescOp.getLoc(),
newDescOp.getResult(),
getPrefetchOffsets(prefetchOffset), readCacheHint,
diff --git a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
index ba259f311d76e..b62188c9762f4 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
@@ -101,3 +101,34 @@ 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 {{prefetch load offset depends on unsupported loop-local values}}
+ %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..51fd3b491aed7 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -433,6 +433,49 @@ 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: %[[PREFETCH_DESC:.*]] = xegpu.create_nd_tdesc %arg0
+ // CHECK-SAME: !xegpu.tensor_desc<128x16xf16>
+ // CHECK: xegpu.prefetch_nd %[[PREFETCH_DESC]]
+ // CHECK: xegpu.prefetch_nd %[[PREFETCH_DESC]]
+ // CHECK: scf.for %[[IV:.*]] = %{{.*}} to %[[K]] step %{{.*}} {
+ // CHECK: %[[FUTURE:.*]] = arith.addi %[[IV]]
+ // CHECK-NOT: scf.if
+ // CHECK: xegpu.prefetch_nd %[[PREFETCH_DESC]][{{.*}}, %[[FUTURE]]]
+ // CHECK: memref.subview %arg0[0, %[[IV]]]
+ // CHECK: xegpu.load_nd
+ 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>
+ %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}) {
+ %func = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %load = transform.structured.match ops{["xegpu.load_nd"]} in %func : (!transform.any_op) -> !transform.any_op
+ %desc = transform.xegpu.insert_prefetch %load nb_prefetch = 2 : (!transform.any_op) -> !transform.any_op
+ %prefetch = transform.get_consumers_of_result %desc[0] : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
// CHECK-LABEL: @convert_layout_a
func.func @convert_layout_a(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
%c0 = arith.constant 0 : index
>From 93c7ac971142f737273f7ef4aba638a83ef603cd Mon Sep 17 00:00:00 2001
From: Charitha Saumya <charitha.saumya.gusthinna.waduge at intel.com>
Date: Thu, 1 Oct 2026 21:42:45 +0000
Subject: [PATCH 2/2] document two cases
---
.../XeGPU/TransformOps/XeGPUTransformOps.td | 9 +-
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 256 +++++++++++-------
.../Dialect/XeGPU/transform-ops-invalid.mlir | 110 +++++++-
mlir/test/Dialect/XeGPU/transform-ops.mlir | 41 ++-
4 files changed, 300 insertions(+), 116 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 1559648c356e3..5ef0c25e003d3 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -154,14 +154,7 @@ def InsertPrefetchOp : Op<Transform_Dialect, "xegpu.insert_prefetch", [
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).
- For a descriptor made from loop-local, rank-preserving, unit-stride
- `memref.subview` ops, the prefetch uses a descriptor of the full source
- memref instead. Subview offsets are added to the prefetch offsets. This
- allows future offsets past the source bounds, as in the static case; the
- descriptor must have `boundary_check = true`. Loop-local offsets may only
- use the loop induction variable directly.
-
- Returns a handle to the created `xegpu.create_nd_tdesc` op.
+ 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 5771dc1a0292d..7ecc3e5e42e0c 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -7,13 +7,15 @@
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.h"
-#include "mlir/Dialect/Arith/Utils/Utils.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>
@@ -350,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,
@@ -401,92 +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);
- Value prefetchSource = descOp.getSource();
- SmallVector<SmallVector<OpFoldResult>> subviewOffsets;
+ 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";
- // Drop rank-preserving, unit-stride subviews made inside the loop. Their
- // offsets move to prefetch_nd, while the descriptor uses the full source
- // memref and remains valid for an out-of-bounds future tile.
- while (Operation *def = prefetchSource.getDefiningOp()) {
- if (!forOp->isAncestor(def))
- break;
- auto subview = dyn_cast<memref::SubViewOp>(def);
- if (!subview ||
- subview.getSourceType().getRank() != subview.getType().getRank() ||
- !areAllConstantIntValue(subview.getMixedStrides(), 1))
- return emitSilenceableFailure(getLoc())
- << "loop-local prefetch requires rank-preserving unit-stride "
- "subviews";
- subviewOffsets.push_back(subview.getMixedOffsets());
- prefetchSource = subview.getSource();
- }
- if (auto arg = dyn_cast<BlockArgument>(prefetchSource);
- arg && arg.getOwner() == forOp.getBody())
+ SmallPtrSet<Operation *, 8> seen;
+ if (failed(collectLoopLocalDescriptorDependancyChain(
+ descOp.getResult(), forOp, seen, descSlice)))
return emitSilenceableFailure(getLoc())
- << "prefetch source depends on a loop iter_arg";
-
- // Only the induction variable itself can be replaced at both the
- // prologue and main-loop prefetch sites.
- auto supportedOffset = [&](Value offset) {
- if (offset == forOp.getInductionVar())
- return true;
- if (Operation *def = offset.getDefiningOp();
- def && forOp->isAncestor(def))
- return false;
- if (auto arg = dyn_cast<BlockArgument>(offset);
- arg && arg.getOwner() == forOp.getBody())
- return false;
- return true;
- };
- for (Value offset : loadOp.getOffsets())
- if (!supportedOffset(offset))
- return emitSilenceableFailure(getLoc())
- << "prefetch load offset depends on unsupported loop-local "
- "values";
- for (const auto &offsets : subviewOffsets)
- for (OpFoldResult offset : offsets)
- if (auto value = dyn_cast<Value>(offset);
- value && !supportedOffset(value))
+ << "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 offset depends on unsupported loop-local "
- "values";
- if (!isa<MemRefType>(prefetchSource.getType()))
- return emitSilenceableFailure(getLoc())
- << "loop-local prefetch requires a memref source";
- if (auto *def = prefetchSource.getDefiningOp();
- def && forOp->isAncestor(def))
+ << "prefetch subview offsets must use the induction variable "
+ "directly or be loop-invariant";
+ }
+ }
+ if (!hasSubview || !subviewUsesIV)
return emitSilenceableFailure(getLoc())
- << "prefetch source depends on unsupported loop-local values";
+ << "loop-local prefetch requires a subview offset that directly "
+ "uses the loop induction variable";
}
- // The prefetch descriptor uses a source that dominates the loop.
+ // 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 =
- loopLocalDesc
- ? xegpu::CreateNdDescOp::create(
- rewriter, descOp.getLoc(), descOp.getType(),
- cast<TypedValue<MemRefType>>(prefetchSource))
- : 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());
@@ -495,29 +556,15 @@ 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);
- if (loopLocalDesc) {
- auto mappedValue = [&](OpFoldResult offset) {
- if (auto value = dyn_cast<Value>(offset))
- return mapping.lookupOrDefault(value);
- return getValueOrCreateConstantIndexOp(rewriter, forOp.getLoc(),
- offset);
- };
- SmallVector<OpFoldResult> offsets;
- for (auto [d, loadOffset] : llvm::enumerate(loadOp.getMixedOffsets())) {
- Value combined = mappedValue(loadOffset);
- for (const auto &subview : subviewOffsets)
- combined = rewriter.createOrFold<arith::AddIOp>(
- forOp.getLoc(), combined, mappedValue(subview[d]));
- offsets.push_back(combined);
- }
- return offsets;
- }
SmallVector<Value> dynamicOffsets =
llvm::map_to_vector(loadOp.getOffsets(), [&](Value v) {
return mapping.lookupOrDefault(v);
@@ -526,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);
+ 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, newDescOp.getLoc(),
- newDescOp.getResult(),
+ 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 b62188c9762f4..886455ff5995b 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
@@ -127,7 +127,115 @@ func.func @insert_prefetch_loop_local_offset(%arg0: memref<128x?xf16>) {
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 offset depends on unsupported loop-local values}}
+ // 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 51fd3b491aed7..c0b00f72d28e2 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -438,23 +438,33 @@ 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
+ // CHECK: %[[K:.*]] = memref.dim %arg0
%k = memref.dim %arg0, %c1 : memref<128x?xf16>
- // CHECK: %[[PREFETCH_DESC:.*]] = xegpu.create_nd_tdesc %arg0
- // CHECK-SAME: !xegpu.tensor_desc<128x16xf16>
- // CHECK: xegpu.prefetch_nd %[[PREFETCH_DESC]]
- // CHECK: xegpu.prefetch_nd %[[PREFETCH_DESC]]
- // CHECK: scf.for %[[IV:.*]] = %{{.*}} to %[[K]] step %{{.*}} {
- // CHECK: %[[FUTURE:.*]] = arith.addi %[[IV]]
- // CHECK-NOT: scf.if
- // CHECK: xegpu.prefetch_nd %[[PREFETCH_DESC]][{{.*}}, %[[FUTURE]]]
- // CHECK: memref.subview %arg0[0, %[[IV]]]
- // CHECK: xegpu.load_nd
+ // 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_END]] step %{{.*}} {
+ // CHECK: %[[INIT_SUB:.*]] = memref.subview %[[OUTER]][0, %[[INIT_IV]]]
+ // CHECK: %[[INIT_DESC:.*]] = xegpu.create_nd_tdesc %[[INIT_SUB]]
+ // CHECK: xegpu.prefetch_nd %[[INIT_DESC]][0, 0]
+ // CHECK: }
+ // CHECK: scf.for %[[IV:.*]] = %{{.*}} to %[[K]] step %{{.*}} {
+ // CHECK: %[[FUTURE:.*]] = arith.addi %[[IV]]
+ // CHECK: %[[IN_BOUNDS:.*]] = arith.cmpi slt, %[[FUTURE]], %[[K]]
+ // CHECK: scf.if %[[IN_BOUNDS]] {
+ // CHECK: %[[MAIN_SUB:.*]] = memref.subview %[[OUTER]][0, %[[FUTURE]]]
+ // CHECK: %[[MAIN_DESC:.*]] = xegpu.create_nd_tdesc %[[MAIN_SUB]]
+ // CHECK: xegpu.prefetch_nd %[[MAIN_DESC]][0, 0]
+ // CHECK: }
+ // CHECK: xegpu.load_nd
+ // CHECK: }
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: ?>>
+ %tile = memref.subview %outer[0, %iv] [128, %size] [1, 1]
+ : memref<128x?xf16, strided<[?, 1]>>
+ to memref<128x?xf16, strided<[?, 1], offset: ?>>
%desc = xegpu.create_nd_tdesc %tile
: memref<128x?xf16, strided<[?, 1], offset: ?>>
-> !xegpu.tensor_desc<128x16xf16>
@@ -469,7 +479,10 @@ module attributes {transform.with_named_sequence} {
%func = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
%load = transform.structured.match ops{["xegpu.load_nd"]} in %func : (!transform.any_op) -> !transform.any_op
%desc = transform.xegpu.insert_prefetch %load nb_prefetch = 2 : (!transform.any_op) -> !transform.any_op
- %prefetch = transform.get_consumers_of_result %desc[0] : (!transform.any_op) -> !transform.any_op
+ transform.foreach %desc : !transform.any_op {
+ ^bb0(%one_desc: !transform.any_op):
+ %prefetch = transform.get_consumers_of_result %one_desc[0] : (!transform.any_op) -> !transform.any_op
+ }
transform.yield
}
}
More information about the Mlir-commits
mailing list