[Mlir-commits] [mlir] [mlir][xegpu] Allow `xegpu.insert_prefetch` to override Load op's tile size. (PR #215419)
Charitha Saumya
llvmlistbot at llvm.org
Fri Aug 14 09:19:58 PDT 2026
https://github.com/charithaintc updated https://github.com/llvm/llvm-project/pull/215419
>From 8b23f4be6c0f9ea6b5464f312f6c43071de3975c Mon Sep 17 00:00:00 2001
From: Charitha Saumya <charitha.saumya.gusthinna.waduge at intel.com>
Date: Mon, 10 Aug 2026 21:35:25 +0000
Subject: [PATCH] allow prefetch tile size to be custom
---
.../XeGPU/TransformOps/XeGPUTransformOps.td | 20 ++-
.../XeGPU/TransformOps/XeGPUTransformOps.cpp | 35 +++++
mlir/python/mlir/dialects/transform/xegpu.py | 19 ++-
.../Dialect/XeGPU/transform-ops-invalid.mlir | 61 ++++++++
mlir/test/Dialect/XeGPU/transform-ops.mlir | 145 ++++++++++++++++++
.../python/dialects/transform_xegpu_ext.py | 40 +++++
6 files changed, 318 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
index 40b9136874e7c..6b8547921c451 100644
--- a/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/TransformOps/XeGPUTransformOps.td
@@ -146,6 +146,7 @@ def SetGPULaunchThreadsOp
}
def InsertPrefetchOp : Op<Transform_Dialect, "xegpu.insert_prefetch", [
+ AttrSizedOperandSegments,
DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
TransformOpInterface
]> {
@@ -156,11 +157,20 @@ 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). Returns a handle
to the created `xegpu.create_nd_desc` op.
+
+ By default the prefetched tile shape is taken from the load op's tensor
+ descriptor. The optional `prefetch_tile_shape` argument overrides it.
+ The prefetch offsets are always derived from the load op's offsets, so
+ the prefetches still start at the loaded tile's origin and advance by
+ the loop step. `prefetch_tile_shape` must have the same rank as
+ the load op's tensor descriptor and all its values must be positive.
}];
let arguments = (ins TransformHandleTypeInterface:$target,
Optional<TransformAnyParamTypeOrAnyHandle>:$dynamic_nb_prefetch,
- DefaultValuedOptionalAttr<I64Attr, "1">:$static_nb_prefetch
+ Variadic<TransformAnyParamTypeOrAnyHandle>:$prefetch_tile_shape,
+ DefaultValuedOptionalAttr<I64Attr, "1">:$static_nb_prefetch,
+ DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$static_prefetch_tile_shape
);
let results = (outs TransformHandleTypeInterface:$desc_op);
@@ -168,6 +178,8 @@ def InsertPrefetchOp : Op<Transform_Dialect, "xegpu.insert_prefetch", [
let assemblyFormat = [{
$target
`nb_prefetch` `=` ($dynamic_nb_prefetch^):($static_nb_prefetch)?
+ (`prefetch_tile_shape` `=` custom<DynamicIndexList>($prefetch_tile_shape,
+ $static_prefetch_tile_shape)^)?
attr-dict `:` functional-type(operands, results)
}];
@@ -184,6 +196,12 @@ def InsertPrefetchOp : Op<Transform_Dialect, "xegpu.insert_prefetch", [
return OpFoldResult(IntegerAttr::get(
IntegerType::get(cxt, 64), getStaticNbPrefetch()));
}
+
+ ::llvm::SmallVector<::mlir::OpFoldResult> getMixedPrefetchTileShape() {
+ Builder b(getContext());
+ return getMixedValues(getStaticPrefetchTileShape(),
+ getPrefetchTileShape(), b);
+ }
}];
}
diff --git a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
index d10265772cacb..a6053ee0c5670 100644
--- a/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/TransformOps/XeGPUTransformOps.cpp
@@ -406,10 +406,44 @@ transform::InsertPrefetchOp::apply(transform::TransformRewriter &rewriter,
return emitSilenceableFailure(getLoc()) << "Could not find descriptor op.";
auto descOp = *maybeDescOp;
+ // Optionally override the prefetched tile shape. By default the shape of the
+ // load op's descriptor is used.
+ xegpu::TensorDescType prefetchDescType = descOp.getType();
+ if (!getMixedPrefetchTileShape().empty()) {
+ SmallVector<int32_t> prefetchTileShape;
+ auto status = convertMixedValuesToInt(state, (*this), prefetchTileShape,
+ getMixedPrefetchTileShape());
+ if (!status.succeeded())
+ return status;
+
+ if (static_cast<int64_t>(prefetchTileShape.size()) !=
+ prefetchDescType.getRank()) {
+ auto diag = emitSilenceableFailure(getLoc())
+ << "prefetch_tile_shape rank (" << prefetchTileShape.size()
+ << ") must match the load op descriptor rank ("
+ << prefetchDescType.getRank() << ")";
+ diag.attachNote(loadOp.getLoc()) << "load op";
+ return diag;
+ }
+ if (llvm::any_of(prefetchTileShape, [](int32_t dim) { return dim <= 0; }))
+ return emitSilenceableFailure(getLoc())
+ << "prefetch_tile_shape values must be positive integers.";
+
+ // Drop the layout as it is not guaranteed to be valid for the new shape.
+ // A layout can be attached to the emitted prefetch ops afterwards.
+ prefetchDescType = xegpu::TensorDescType::get(
+ getContext(),
+ SmallVector<int64_t>(prefetchTileShape.begin(),
+ prefetchTileShape.end()),
+ prefetchDescType.getElementType(), prefetchDescType.getEncoding(),
+ /*layout=*/nullptr);
+ }
+
// Clone desc op outside the loop.
rewriter.setInsertionPoint(forOp);
auto newDescOp =
cast<xegpu::CreateNdDescOp>(rewriter.clone(*descOp.getOperation()));
+ newDescOp.getResult().setType(prefetchDescType);
// Clone reduction loop to emit initial prefetches.
// Compute upper bound of the init loop: start + nbPrefetch * step.
@@ -474,6 +508,7 @@ void transform::InsertPrefetchOp::getEffects(
::llvm::SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
onlyReadsHandle(getTargetMutable(), effects);
onlyReadsHandle(getDynamicNbPrefetchMutable(), effects);
+ onlyReadsHandle(getPrefetchTileShapeMutable(), effects);
producesHandle(getOperation()->getOpResults(), effects);
modifiesPayload(effects);
}
diff --git a/mlir/python/mlir/dialects/transform/xegpu.py b/mlir/python/mlir/dialects/transform/xegpu.py
index 6e27e5c8ecfa6..6eafe7f75041a 100644
--- a/mlir/python/mlir/dialects/transform/xegpu.py
+++ b/mlir/python/mlir/dialects/transform/xegpu.py
@@ -171,6 +171,7 @@ def __init__(
target: Union[Operation, Value],
*,
nb_prefetch: Optional[MixedInt] = 1,
+ prefetch_tile_shape: Optional[MixedValues] = None,
loc=None,
ip=None,
):
@@ -183,11 +184,20 @@ def __init__(
elif isinstance(nb_prefetch, (Operation, Value, OpView)):
dynamic_nb_prefetch = nb_prefetch
+ prefetch_tile_shape = [] if prefetch_tile_shape is None else prefetch_tile_shape
+ (
+ dynamic_prefetch_tile_shape,
+ static_prefetch_tile_shape,
+ _,
+ ) = _dispatch_dynamic_index_list(prefetch_tile_shape)
+
super().__init__(
transform.AnyOpType.get(),
target,
+ prefetch_tile_shape=dynamic_prefetch_tile_shape,
dynamic_nb_prefetch=dynamic_nb_prefetch,
static_nb_prefetch=static_nb_prefetch,
+ static_prefetch_tile_shape=static_prefetch_tile_shape,
loc=loc,
ip=ip,
)
@@ -197,10 +207,17 @@ def insert_prefetch(
target: Union[Operation, Value],
*,
nb_prefetch: Optional[MixedInt] = 1,
+ prefetch_tile_shape: Optional[MixedValues] = None,
loc=None,
ip=None,
) -> OpResult:
- return InsertPrefetchOp(target, nb_prefetch=nb_prefetch, loc=loc, ip=ip).result
+ return InsertPrefetchOp(
+ target,
+ nb_prefetch=nb_prefetch,
+ prefetch_tile_shape=prefetch_tile_shape,
+ loc=loc,
+ ip=ip,
+ ).result
@_ods_cext.register_operation(_Dialect, replace=True)
diff --git a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
index ba259f311d76e..2166f7ce0248c 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops-invalid.mlir
@@ -101,3 +101,64 @@ module attributes {transform.with_named_sequence} {
transform.yield
}
}
+
+// -----
+
+func.func @insert_prefetch_bad_tile_shape_rank(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
+ %c32 = arith.constant 32 : index
+ %c4096 = arith.constant 4096 : index
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
+ %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+ %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>
+ %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {
+ // expected-note at below {{load op}}
+ %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>
+ %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
+ scf.yield %7 : vector<256x256xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
+ %2 = transform.xegpu.get_load_op %1 : (!transform.any_value) -> !transform.any_op
+ // expected-error at below {{prefetch_tile_shape rank (3) must match the load op descriptor rank (2)}}
+ %3 = transform.xegpu.insert_prefetch %2 nb_prefetch = 1 prefetch_tile_shape = [64, 32, 8] : (!transform.any_op) -> !transform.any_op
+ transform.yield
+ }
+}
+
+// -----
+
+func.func @insert_prefetch_non_positive_tile_shape(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
+ %c32 = arith.constant 32 : index
+ %c4096 = arith.constant 4096 : index
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
+ %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+ %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>
+ %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {
+ %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>
+ %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
+ scf.yield %7 : vector<256x256xf16>
+ }
+ return
+}
+
+module attributes {transform.with_named_sequence} {
+ transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+ %0 = transform.structured.match ops{["xegpu.dpas"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
+ %2 = transform.xegpu.get_load_op %1 : (!transform.any_value) -> !transform.any_op
+ // expected-error at below {{prefetch_tile_shape values must be positive integers.}}
+ %3 = transform.xegpu.insert_prefetch %2 nb_prefetch = 1 prefetch_tile_shape = [64, 0] : (!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 3daa74d223946..65b45a1149993 100644
--- a/mlir/test/Dialect/XeGPU/transform-ops.mlir
+++ b/mlir/test/Dialect/XeGPU/transform-ops.mlir
@@ -433,6 +433,151 @@ module attributes {transform.with_named_sequence} {
// -----
+// Prefetch tile shape overrides the load op tile shape, but prefetch offsets
+// still follow the load op offsets.
+// CHECK-LABEL: @insert_prefetch_tile_shape
+func.func @insert_prefetch_tile_shape(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
+ // CHECK: %[[C64:.+]] = arith.constant 64 : index
+ // CHECK: %[[C32:.+]] = arith.constant 32 : index
+ %c32 = arith.constant 32 : index
+ %c4096 = arith.constant 4096 : index
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
+ // CHECK: xegpu.create_nd_tdesc %arg0
+ // CHECK: xegpu.create_nd_tdesc %arg1
+ // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0
+ // CHECK-SAME: !xegpu.tensor_desc<64x32xf16
+ // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[C0]]]
+ // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[C32]]]
+ %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+ %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>
+ // CHECK: scf.for %[[ARG3:.+]] = %[[C0]]
+ %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {
+ // CHECK: %[[ADD:.+]] = arith.addi %[[ARG3]], %[[C64]]
+ // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[ADD]]]
+ // CHECK: xegpu.load_nd {{.*}} !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>
+ %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
+ scf.yield %7 : vector<256x256xf16>
+ }
+ 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
+ %0 = transform.structured.match ops{["xegpu.dpas"]} in %func : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
+ %2 = transform.xegpu.get_load_op %1 : (!transform.any_value) -> !transform.any_op
+ // CHECK: transform.xegpu.insert_prefetch %{{.*}} nb_prefetch = 2 prefetch_tile_shape = [64, 32]
+ %3 = transform.xegpu.insert_prefetch %2 nb_prefetch = 2 prefetch_tile_shape = [64, 32] : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.canonicalization
+ } : !transform.any_op
+
+ transform.yield
+ }
+}
+
+// -----
+
+// Prefetch tile shape can be given as transform params.
+// CHECK-LABEL: @insert_prefetch_tile_shape_param
+func.func @insert_prefetch_tile_shape_param(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
+ // CHECK: %[[C32:.+]] = arith.constant 32 : index
+ %c32 = arith.constant 32 : index
+ %c4096 = arith.constant 4096 : index
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
+ // CHECK: xegpu.create_nd_tdesc %arg0
+ // CHECK: xegpu.create_nd_tdesc %arg1
+ // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg1
+ // CHECK-SAME: !xegpu.tensor_desc<16x128xf16
+ // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[C0]]]
+ %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16>
+ %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>
+ // CHECK: scf.for %[[ARG3:.+]] = %[[C0]]
+ %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {
+ // CHECK: %[[ADD:.+]] = arith.addi %[[ARG3]], %[[C32]]
+ // CHECK: xegpu.prefetch_nd %[[V0]][%[[ADD]], %[[C0]]]
+ %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16> -> vector<256x32xf16>
+ %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>
+ %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
+ scf.yield %7 : vector<256x256xf16>
+ }
+ 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
+ %0 = transform.structured.match ops{["xegpu.dpas"]} in %func : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_operand %0[1] : (!transform.any_op) -> !transform.any_value
+ %2 = transform.xegpu.get_load_op %1 : (!transform.any_value) -> !transform.any_op
+ %k = transform.param.constant 16 : i64 -> !transform.param<i64>
+ // CHECK: transform.xegpu.insert_prefetch %{{.*}} nb_prefetch = 1 prefetch_tile_shape = [%{{.*}}, 128]
+ %3 = transform.xegpu.insert_prefetch %2 nb_prefetch = 1 prefetch_tile_shape = [%k, 128] : (!transform.any_op, !transform.param<i64>) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.canonicalization
+ } : !transform.any_op
+
+ transform.yield
+ }
+}
+
+// -----
+
+// The descriptor layout is dropped when the prefetch tile shape is overridden,
+// since it is not guaranteed to be valid for the new shape.
+// CHECK-LABEL: @insert_prefetch_tile_shape_drops_layout
+func.func @insert_prefetch_tile_shape_drops_layout(%arg0: memref<4096x4096xf16>, %arg1: memref<4096x4096xf16>, %arg2: memref<4096x4096xf16>) {
+ %c32 = arith.constant 32 : index
+ %c4096 = arith.constant 4096 : index
+ // CHECK: %[[C0:.+]] = arith.constant 0 : index
+ %c0 = arith.constant 0 : index
+ %0 = xegpu.create_nd_tdesc %arg2 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x256xf16>
+ %1 = xegpu.load_nd %0[%c0, %c0] : !xegpu.tensor_desc<256x256xf16> -> vector<256x256xf16>
+ // CHECK: xegpu.create_nd_tdesc %arg0
+ // CHECK-SAME: #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 8]>
+ // CHECK: xegpu.create_nd_tdesc %arg1
+ // CHECK: %[[V0:.+]] = xegpu.create_nd_tdesc %arg0
+ // CHECK-SAME: !xegpu.tensor_desc<64x32xf16>
+ // CHECK: xegpu.prefetch_nd %[[V0]][%[[C0]], %[[C0]]]
+ // CHECK-SAME: !xegpu.tensor_desc<64x32xf16>
+ %3 = xegpu.create_nd_tdesc %arg0 : memref<4096x4096xf16> -> !xegpu.tensor_desc<256x32xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 8]>>
+ %4 = xegpu.create_nd_tdesc %arg1 : memref<4096x4096xf16> -> !xegpu.tensor_desc<32x256xf16>
+ %2 = scf.for %arg3 = %c0 to %c4096 step %c32 iter_args(%arg4 = %1) -> (vector<256x256xf16>) {
+ %5 = xegpu.load_nd %3[%c0, %arg3] : !xegpu.tensor_desc<256x32xf16, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 8]>> -> vector<256x32xf16>
+ %6 = xegpu.load_nd %4[%arg3, %c0] : !xegpu.tensor_desc<32x256xf16> -> vector<32x256xf16>
+ %7 = xegpu.dpas %5, %6, %arg4 : vector<256x32xf16>, vector<32x256xf16>, vector<256x256xf16> -> vector<256x256xf16>
+ scf.yield %7 : vector<256x256xf16>
+ }
+ 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
+ %0 = transform.structured.match ops{["xegpu.dpas"]} in %func : (!transform.any_op) -> !transform.any_op
+ %1 = transform.get_operand %0[0] : (!transform.any_op) -> !transform.any_value
+ %2 = transform.xegpu.get_load_op %1 : (!transform.any_value) -> !transform.any_op
+ // CHECK: transform.xegpu.insert_prefetch %{{.*}}
+ %3 = transform.xegpu.insert_prefetch %2 nb_prefetch = 1 prefetch_tile_shape = [64, 32] : (!transform.any_op) -> !transform.any_op
+ transform.apply_patterns to %func {
+ transform.apply_patterns.canonicalization
+ } : !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
diff --git a/mlir/test/python/dialects/transform_xegpu_ext.py b/mlir/test/python/dialects/transform_xegpu_ext.py
index 5d5db1919af14..353c0ab0a16d7 100644
--- a/mlir/test/python/dialects/transform_xegpu_ext.py
+++ b/mlir/test/python/dialects/transform_xegpu_ext.py
@@ -194,6 +194,46 @@ def insertPrefetchNbPrefetchParam():
# CHECK-SAME: nb_prefetch = %[[PARAM_OP]]
+ at run
+def insertPrefetchTileShape():
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.OperationType.get("xegpu.load_nd"),
+ )
+ with InsertionPoint(sequence.body):
+ xegpu.insert_prefetch(
+ sequence.bodyTarget, nb_prefetch=2, prefetch_tile_shape=[64, 32]
+ )
+ transform.YieldOp()
+ # CHECK-LABEL: TEST: insertPrefetchTileShape
+ # CHECK: transform.xegpu.insert_prefetch
+ # CHECK-SAME: nb_prefetch = 2
+ # CHECK-SAME: prefetch_tile_shape = [64, 32]
+
+
+ at run
+def insertPrefetchTileShapeParam():
+ sequence = transform.SequenceOp(
+ transform.FailurePropagationMode.Propagate,
+ [],
+ transform.OperationType.get("xegpu.load_nd"),
+ )
+ with InsertionPoint(sequence.body):
+ int32_t = IntegerType.get_signless(32)
+ param_int32_t = transform.ParamType.get(int32_t)
+ dim_param = transform.ParamConstantOp(
+ param_int32_t,
+ IntegerAttr.get(int32_t, 64),
+ )
+ xegpu.insert_prefetch(sequence.bodyTarget, prefetch_tile_shape=[dim_param, 32])
+ transform.YieldOp()
+ # CHECK-LABEL: TEST: insertPrefetchTileShapeParam
+ # CHECK: %[[PARAM_OP:.*]] = transform.param.constant 64
+ # CHECK: transform.xegpu.insert_prefetch
+ # CHECK-SAME: prefetch_tile_shape = [%[[PARAM_OP]], 32]
+
+
@run
def ConvertLayoutMinimal():
sequence = transform.SequenceOp(
More information about the Mlir-commits
mailing list