[Mlir-commits] [mlir] e3ba084 - [mlir][xegpu] Fix high-D nd base_height for non-packed plane strides (#225556)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Oct 2 09:03:34 PDT 2026
Author: Jianhui Li
Date: 2026-10-02T09:03:16-07:00
New Revision: e3ba0841a01eef468062f736e59ca7d014a4f11f
URL: https://github.com/llvm/llvm-project/commit/e3ba0841a01eef468062f736e59ca7d014a4f11f
DIFF: https://github.com/llvm/llvm-project/commit/e3ba0841a01eef468062f736e59ca7d014a4f11f.diff
LOG: [mlir][xegpu] Fix high-D nd base_height for non-packed plane strides (#225556)
CreateNdDescToXeVMPattern lowers a >2D nd descriptor by viewing the
source as one flattened 2D block surface. It set the surface height to
the product of all dims but the innermost, which is the right row extent
only when the leading (batch) planes are packed back to back. A subview
of a wider source leaves gaps between them, and the last planes end up
outside base_height, where the HW boundary check silently drops every
access to them.
This PR fixes the issue by using the following formula:
base_height = size[R-2] + sum_d (size[d] - 1) * (stride[d] /
stride[R-2])
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply at anthropic.com>
Added:
mlir/test/Integration/Dialect/XeGPU/WG/store_nd_3d_subview.mlir
Modified:
mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
index e1a82e850a2731..b2a81f01887bb9 100644
--- a/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
+++ b/mlir/lib/Conversion/XeGPUToXeVM/XeGPUToXeVM.cpp
@@ -214,27 +214,32 @@ translateStoreXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
//
// High-D (>2D) nd descriptors are lowered by viewing the source as a single
-// flattened 2D plane: `base_height` is the product of all dims but the
-// innermost, so the 2D-block surface covers every leading (batch) plane at
-// once, and a batch position becomes a row offset into it. Leaving `base_ptr`
-// at the true base means an out-of-range batch index lands past `base_height`,
-// where the HW boundary check handles it, instead of aiming the surface at
-// unmapped memory. Encoding the leading strides as row counts (`stride[d] /
-// stride[R-2]`) also makes them dimensionless, so they survive element-type
-// repacking (e.g. the f16 -> i32 transpose repack) without a unit conversion.
+// flattened 2D plane, so the 2D-block surface covers every leading (batch)
+// plane at once and a batch position becomes a row offset into it. The surface
+// height is the row extent the leading dims reach:
+// base_height = size[R-2] + sum_d (size[d] - 1) * (stride[d] / stride[R-2])
+//
+// Leaving `base_ptr` at the true base means an out-of-range batch index lands
+// past `base_height`, where the HW boundary check handles it, instead of aiming
+// the surface at unmapped memory. Encoding the leading strides as row counts
+// (`stride[d] / stride[R-2]`) also makes them dimensionless, so they survive
+// element-type repacking (e.g. the f16 -> i32 transpose repack) without a unit
+// conversion.
//
// Limitations of the flattened-plane view:
// 1. Each leading stride must be a whole number of rows. A source with gaps
// between planes (`stride[d] % stride[R-2] != 0`) is not lowered. This can
// only be checked when the strides are static; for dynamic strides the
// divisibility is assumed.
-// 2. `base_height` grows to the product of the leading dims, so a source with
-// a large batch x head x sequence extent can exceed the HW 2D-block
-// surface height.
-// 3. Plane boundaries are invisible to the boundary check: a tile whose rows
-// run past `size[R-2]` reads the next plane's rows instead of the zero
-// padding a per-plane surface would return. This only matters when
-// `size[R-2]` is not a multiple of the tile height.
+// 2. `base_height` grows with the leading extent, so a source with a large
+// batch x head x sequence extent can exceed the HW 2D-block surface height.
+// 3. Plane boundaries are invisible to the boundary check, which only knows
+// the surface: a tile whose rows run past `size[R-2]` spills into the
+// following rows of the surface -- the next plane, or the padding between
+// planes -- rather than being clipped there as it would be on a per-plane
+// surface. A load reads those rows instead of returning zeros, and a store
+// overwrites them, so it can corrupt the next plane. This only matters
+// when `size[R-2]` is not a multiple of the tile height.
//
class CreateNdDescToXeVMPattern
@@ -347,15 +352,37 @@ class CreateNdDescToXeVMPattern
};
// The descriptor's innermost 2 dims are the 2D tile (H, W).
Value baseShapeW = createOffset(mixedSizes, rank - 1);
+ // Pitch is the stride of dim rank-2 (the row stride of the 2D tile).
+ Value basePitch = createOffset(mixedStrides, rank - 2);
+ // Leading (batch) strides as a number of rows (`stride[d] / pitch`).
+ SmallVector<Value> leadingRowStrides;
+ for (int64_t d = 0; d < rank - 2; ++d) {
+ std::optional<int64_t> leading = getConstantIntValue(mixedStrides[d]);
+ std::optional<int64_t> pitch =
+ getConstantIntValue(mixedStrides[rank - 2]);
+ if (leading && pitch && *pitch != 0)
+ leadingRowStrides.push_back(arith::ConstantIntOp::create(
+ rewriter, loc, payloadElemTy, *leading / *pitch));
+ else
+ leadingRowStrides.push_back(arith::DivUIOp::create(
+ rewriter, loc, createOffset(mixedStrides, d), basePitch));
+ }
// Height of the flattened plane: every leading (batch) plane is stacked
// into the surface, so the boundary check covers an out-of-range batch.
+ // base_height = size[R-2] + sum_d (size[d] - 1) * leadingRowStride[d]
// For rank 2 this is just size[0].
Value baseShapeH = createOffset(mixedSizes, rank - 2);
- for (int64_t d = 0; d < rank - 2; ++d)
- baseShapeH = arith::MulIOp::create(rewriter, loc, baseShapeH,
- createOffset(mixedSizes, d));
- // Pitch is the stride of dim rank-2 (the row stride of the 2D tile).
- Value basePitch = createOffset(mixedStrides, rank - 2);
+ if (rank > 2) {
+ Value one = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, 1);
+ for (int64_t d = 0; d < rank - 2; ++d) {
+ Value planesBelow = rewriter.createOrFold<arith::SubIOp>(
+ loc, createOffset(mixedSizes, d), one);
+ Value rows = rewriter.createOrFold<arith::MulIOp>(loc, planesBelow,
+ leadingRowStrides[d]);
+ baseShapeH =
+ rewriter.createOrFold<arith::AddIOp>(loc, baseShapeH, rows);
+ }
+ }
// Populate payload.
Value payLoadAsI64 =
vector::BitCastOp::create(rewriter, loc, payloadI64Ty, payload);
@@ -372,26 +399,13 @@ class CreateNdDescToXeVMPattern
payload =
vector::InsertOp::create(rewriter, loc, basePitch, payload,
static_cast<int>(NdTdescOffset::BasePitch));
- // Leading (batch) strides go into the spare payload slots as a number of
- // rows; the load/store/prefetch lowering turns the batch offsets into a row
- // offset with them. Row units keep them independent of the element type, so
- // an element-type repack cannot put them out of step with the pitch.
- for (int64_t d = 0; d < rank - 2; ++d) {
- std::optional<int64_t> leading = getConstantIntValue(mixedStrides[d]);
- std::optional<int64_t> pitch =
- getConstantIntValue(mixedStrides[rank - 2]);
- Value leadingRowStride;
- if (leading && pitch && *pitch != 0) {
- leadingRowStride = arith::ConstantIntOp::create(
- rewriter, loc, payloadElemTy, *leading / *pitch);
- } else {
- leadingRowStride = arith::DivUIOp::create(
- rewriter, loc, createOffset(mixedStrides, d), basePitch);
- }
+ // The leading (batch) row strides go into the spare payload slots; the
+ // load/store/prefetch lowering turns the batch offsets into a row offset
+ // with them.
+ for (int64_t d = 0; d < rank - 2; ++d)
payload = vector::InsertOp::create(
- rewriter, loc, leadingRowStride, payload,
+ rewriter, loc, leadingRowStrides[d], payload,
static_cast<int>(NdTdescOffset::LeadingStride0) + d);
- }
rewriter.replaceOp(op, payload);
return success();
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
index 939adc951ac385..9530ff9c5ba4de 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/create_nd_tdesc.mlir
@@ -13,8 +13,8 @@ gpu.module @create_nd_tdesc {
// CHECK: %[[BASE_ADDR:.*]] = arith.index_castui %[[VAR0]] : index to i64
// CHECK: %[[CST:.*]] = arith.constant dense<0> : vector<8xi32>
// CHECK: %[[SHAPE_W:.*]] = arith.index_cast %[[ARG3]] : index to i32
- // CHECK: %[[SHAPE_H:.*]] = arith.index_cast %[[ARG2]] : index to i32
// CHECK: %[[PITCH:.*]] = arith.index_cast %[[ARG4]] : index to i32
+ // CHECK: %[[SHAPE_H:.*]] = arith.index_cast %[[ARG2]] : index to i32
// CHECK: %[[VAR6:.*]] = vector.bitcast %[[CST]] : vector<8xi32> to vector<4xi64>
// CHECK: %[[VAR7:.*]] = vector.insert %[[BASE_ADDR]], %[[VAR6]] [0] : i64 into vector<4xi64>
// CHECK: %[[VAR8:.*]] = vector.bitcast %[[VAR7]] : vector<4xi64> to vector<8xi32>
@@ -32,10 +32,10 @@ gpu.module @create_nd_tdesc {
// CHECK: %[[CST_1:.*]] = arith.constant dense<0> : vector<8xi32>
// CHECK: %[[C32_I64:.*]] = arith.constant 32 : i64
// CHECK: %[[SHAPE_W2:.*]] = arith.trunci %[[C32_I64]] : i64 to i32
- // CHECK: %[[C16_I64:.*]] = arith.constant 16 : i64
- // CHECK: %[[SHAPE_H2:.*]] = arith.trunci %[[C16_I64]] : i64 to i32
// CHECK: %[[C32_I64_2:.*]] = arith.constant 32 : i64
// CHECK: %[[PITCH2:.*]] = arith.trunci %[[C32_I64_2]] : i64 to i32
+ // CHECK: %[[C16_I64:.*]] = arith.constant 16 : i64
+ // CHECK: %[[SHAPE_H2:.*]] = arith.trunci %[[C16_I64]] : i64 to i32
// CHECK: %[[VAR14:.*]] = vector.bitcast %[[CST_1]] : vector<8xi32> to vector<4xi64>
// CHECK: %[[VAR15:.*]] = vector.insert %[[BASE_ADDR2_OFFSET:.*]], %[[VAR14]] [0] : i64 into vector<4xi64>
// CHECK: %[[VAR16:.*]] = vector.bitcast %[[VAR15]] : vector<4xi64> to vector<8xi32>
@@ -48,8 +48,8 @@ gpu.module @create_nd_tdesc {
// CHECK: %{{.*}}, %{{.*}}, %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[DYN]] : memref<?x?xf16>
// CHECK: %[[CST_3:.*]] = arith.constant dense<0> : vector<8xi32>
// CHECK: %[[SHAPE_W3:.*]] = arith.index_cast %[[SIZES]]#1 : index to i32
- // CHECK: %[[SHAPE_H3:.*]] = arith.index_cast %[[SIZES]]#0 : index to i32
// CHECK: %[[PITCH3:.*]] = arith.index_cast %[[STRIDES]]#0 : index to i32
+ // CHECK: %[[SHAPE_H3:.*]] = arith.index_cast %[[SIZES]]#0 : index to i32
// CHECK: %[[VAR25:.*]] = vector.bitcast %[[CST_3]] : vector<8xi32> to vector<4xi64>
// CHECK: %[[VAR26:.*]] = vector.insert %{{.*}}, %[[VAR25]] [0] : i64 into vector<4xi64>
// CHECK: %[[VAR27:.*]] = vector.bitcast %[[VAR26]] : vector<4xi64> to vector<8xi32>
@@ -60,24 +60,97 @@ gpu.module @create_nd_tdesc {
gpu.return
}
- // Batched (>2D): base_height spans all planes; slot 5 = batch row stride.
+ // Batched (>2D) with two leading dims, fully dynamic sizes and strides:
+ // base_height is the row extent the batch dims reach,
+ // size[2] + (size[0] - 1) * rows0 + (size[1] - 1) * rows1
+ // with rows_d = stride[d] / pitch; slots 5 and 6 are the batch row strides.
// CHECK-LABEL: gpu.func @create_nd_tdesc_batch_dyn(
- // CHECK-SAME: %[[SRC:.+]]: memref<?x?x?xf16>
- gpu.func @create_nd_tdesc_batch_dyn(%src: memref<?x?x?xf16>) -> vector<8xi32> {
- // CHECK: %{{.+}}, %{{.+}}, %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
- // CHECK: %[[W:.+]] = arith.index_cast %[[SIZES]]#2 : index to i32
- // CHECK: %[[H:.+]] = arith.index_cast %[[SIZES]]#1 : index to i32
- // CHECK: %[[BATCH:.+]] = arith.index_cast %[[SIZES]]#0 : index to i32
- // CHECK: %[[FLAT_H:.+]] = arith.muli %[[H]], %[[BATCH]] : i32
- // CHECK: %[[PITCH:.+]] = arith.index_cast %[[STRIDES]]#1 : index to i32
+ // CHECK-SAME: %[[SRC:.+]]: memref<?x?x?x?xf16>
+ gpu.func @create_nd_tdesc_batch_dyn(%src: memref<?x?x?x?xf16>) -> vector<8xi32> {
+ // CHECK: %{{.+}}, %{{.+}}, %[[SIZES:.+]]:4, %[[STRIDES:.+]]:4 = memref.extract_strided_metadata %[[SRC]]
+ // CHECK: %[[W:.+]] = arith.index_cast %[[SIZES]]#3 : index to i32
+ // CHECK: %[[PITCH:.+]] = arith.index_cast %[[STRIDES]]#2 : index to i32
+ // CHECK: %[[LS0:.+]] = arith.index_cast %[[STRIDES]]#0 : index to i32
+ // CHECK: %[[ROWS0:.+]] = arith.divui %[[LS0]], %[[PITCH]] : i32
+ // CHECK: %[[LS1:.+]] = arith.index_cast %[[STRIDES]]#1 : index to i32
+ // CHECK: %[[ROWS1:.+]] = arith.divui %[[LS1]], %[[PITCH]] : i32
+ // CHECK: %[[H:.+]] = arith.index_cast %[[SIZES]]#2 : index to i32
+ // CHECK: %[[C1I32:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BATCH0:.+]] = arith.index_cast %[[SIZES]]#0 : index to i32
+ // CHECK: %[[BATCH0M1:.+]] = arith.subi %[[BATCH0]], %[[C1I32]] : i32
+ // CHECK: %[[BATCH0_ROWS:.+]] = arith.muli %[[BATCH0M1]], %[[ROWS0]] : i32
+ // CHECK: %[[FLAT_H0:.+]] = arith.addi %[[H]], %[[BATCH0_ROWS]] : i32
+ // CHECK: %[[BATCH1:.+]] = arith.index_cast %[[SIZES]]#1 : index to i32
+ // CHECK: %[[BATCH1M1:.+]] = arith.subi %[[BATCH1]], %[[C1I32]] : i32
+ // CHECK: %[[BATCH1_ROWS:.+]] = arith.muli %[[BATCH1M1]], %[[ROWS1]] : i32
+ // CHECK: %[[FLAT_H:.+]] = arith.addi %[[FLAT_H0]], %[[BATCH1_ROWS]] : i32
// CHECK: %[[P2:.+]] = vector.insert %[[W]], %{{.+}} [2] : i32 into vector<8xi32>
// CHECK: %[[P3:.+]] = vector.insert %[[FLAT_H]], %[[P2]] [3] : i32 into vector<8xi32>
// CHECK: %[[P4:.+]] = vector.insert %[[PITCH]], %[[P3]] [4] : i32 into vector<8xi32>
- // CHECK: %[[LS0:.+]] = arith.index_cast %[[STRIDES]]#0 : index to i32
- // CHECK: %[[ROWS0:.+]] = arith.divui %[[LS0]], %[[PITCH]] : i32
- // CHECK: vector.insert %[[ROWS0]], %[[P4]] [5] : i32 into vector<8xi32>
- %t = xegpu.create_nd_tdesc %src : memref<?x?x?xf16> -> !xegpu.tensor_desc<1x8x16xf16>
- %c = builtin.unrealized_conversion_cast %t : !xegpu.tensor_desc<1x8x16xf16> to vector<8xi32>
+ // CHECK: %[[P5:.+]] = vector.insert %[[ROWS0]], %[[P4]] [5] : i32 into vector<8xi32>
+ // CHECK: vector.insert %[[ROWS1]], %[[P5]] [6] : i32 into vector<8xi32>
+ %t = xegpu.create_nd_tdesc %src : memref<?x?x?x?xf16> -> !xegpu.tensor_desc<1x1x8x16xf16>
+ %c = builtin.unrealized_conversion_cast %t : !xegpu.tensor_desc<1x1x8x16xf16> to vector<8xi32>
+ gpu.return %c : vector<8xi32>
+ }
+
+ // Batched (>2D) with the maximum three leading dims, over a subview of a
+ // wider source taken at a nonzero offset. The innermost planes are 64 rows
+ // apart but only 32 rows tall, so the planes are not packed back to back.
+ // The subview's offset belongs to base_ptr, so base_height stays the row
+ // extent measured from that shifted base:
+ // 32 + 1 * 768 + 1 * 256 + 3 * 64 = 1248
+ // and the furthest row the batch dims reach is 768 + 256 + 3 * 64 + 31 =
+ // 1247, just inside it. All three spare payload slots carry a batch row
+ // stride.
+ // CHECK-LABEL: gpu.func @create_nd_tdesc_batch_subview(
+ // CHECK-SAME: %[[SRC:.+]]: memref<3x3x4x64x64xf32>
+ gpu.func @create_nd_tdesc_batch_subview(%src: memref<3x3x4x64x64xf32>) -> vector<8xi32> {
+ // The subview offset 1 * 49152 + 1 * 16384 + 16 * 64 = 66560 folds into
+ // base_ptr and leaves the shape fields alone.
+ // CHECK: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %{{.+}}
+ // CHECK: %[[OFF:.+]] = arith.constant 66560 : index
+ // CHECK: %[[PTR_I64:.+]] = arith.index_castui %[[INTPTR]] : index to i64
+ // CHECK: %[[OFF_I64:.+]] = arith.index_castui %[[OFF]] : index to i64
+ // CHECK: %[[ELEM_SZ:.+]] = arith.constant 4 : i64
+ // CHECK: %[[OFF_BYTES:.+]] = arith.muli %[[OFF_I64]], %[[ELEM_SZ]] : i64
+ // CHECK: %[[BASE_PTR:.+]] = arith.addi %[[PTR_I64]], %[[OFF_BYTES]] : i64
+ // W = size[4] = 32, pitch = stride[3] = 64.
+ // CHECK: %[[W:.+]] = arith.trunci %{{.+}} : i64 to i32
+ // CHECK: %[[PITCH:.+]] = arith.trunci %{{.+}} : i64 to i32
+ // Batch row strides = stride[d] / pitch: 49152 / 64, 16384 / 64, 4096 / 64.
+ // CHECK: %[[ROWS0:.+]] = arith.constant 768 : i32
+ // CHECK: %[[ROWS1:.+]] = arith.constant 256 : i32
+ // CHECK: %[[ROWS2:.+]] = arith.constant 64 : i32
+ // base_height = size[3] + sum_d (size[d] - 1) * rows_d.
+ // CHECK: %[[H:.+]] = arith.trunci %{{.+}} : i64 to i32
+ // CHECK: %[[C1I32:.+]] = arith.constant 1 : i32
+ // CHECK: %[[BATCH0:.+]] = arith.trunci %{{.+}} : i64 to i32
+ // CHECK: %[[BATCH0M1:.+]] = arith.subi %[[BATCH0]], %[[C1I32]] : i32
+ // CHECK: %[[BATCH0_ROWS:.+]] = arith.muli %[[BATCH0M1]], %[[ROWS0]] : i32
+ // CHECK: %[[FLAT_H0:.+]] = arith.addi %[[H]], %[[BATCH0_ROWS]] : i32
+ // CHECK: %[[BATCH1:.+]] = arith.trunci %{{.+}} : i64 to i32
+ // CHECK: %[[BATCH1M1:.+]] = arith.subi %[[BATCH1]], %[[C1I32]] : i32
+ // CHECK: %[[BATCH1_ROWS:.+]] = arith.muli %[[BATCH1M1]], %[[ROWS1]] : i32
+ // CHECK: %[[FLAT_H1:.+]] = arith.addi %[[FLAT_H0]], %[[BATCH1_ROWS]] : i32
+ // CHECK: %[[BATCH2:.+]] = arith.trunci %{{.+}} : i64 to i32
+ // CHECK: %[[BATCH2M1:.+]] = arith.subi %[[BATCH2]], %[[C1I32]] : i32
+ // CHECK: %[[BATCH2_ROWS:.+]] = arith.muli %[[BATCH2M1]], %[[ROWS2]] : i32
+ // CHECK: %[[FLAT_H:.+]] = arith.addi %[[FLAT_H1]], %[[BATCH2_ROWS]] : i32
+ // CHECK: %[[PI64:.+]] = vector.insert %[[BASE_PTR]], %{{.+}} [0] : i64 into vector<4xi64>
+ // CHECK: %[[P2:.+]] = vector.insert %[[W]], %{{.+}} [2] : i32 into vector<8xi32>
+ // CHECK: %[[P3:.+]] = vector.insert %[[FLAT_H]], %[[P2]] [3] : i32 into vector<8xi32>
+ // CHECK: %[[P4:.+]] = vector.insert %[[PITCH]], %[[P3]] [4] : i32 into vector<8xi32>
+ // CHECK: %[[P5:.+]] = vector.insert %[[ROWS0]], %[[P4]] [5] : i32 into vector<8xi32>
+ // CHECK: %[[P6:.+]] = vector.insert %[[ROWS1]], %[[P5]] [6] : i32 into vector<8xi32>
+ // CHECK: vector.insert %[[ROWS2]], %[[P6]] [7] : i32 into vector<8xi32>
+ %sub = memref.subview %src[1, 1, 0, 16, 0] [2, 2, 4, 32, 32] [1, 1, 1, 1, 1]
+ : memref<3x3x4x64x64xf32>
+ to memref<2x2x4x32x32xf32, strided<[49152, 16384, 4096, 64, 1], offset: 66560>>
+ %t = xegpu.create_nd_tdesc %sub
+ : memref<2x2x4x32x32xf32, strided<[49152, 16384, 4096, 64, 1], offset: 66560>>
+ -> !xegpu.tensor_desc<1x1x1x8x16xf32>
+ %c = builtin.unrealized_conversion_cast %t : !xegpu.tensor_desc<1x1x1x8x16xf32> to vector<8xi32>
gpu.return %c : vector<8xi32>
}
}
diff --git a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
index 408950c0d8f0f4..71053e341aa0c3 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/loadstore_nd.mlir
@@ -99,15 +99,19 @@ gpu.module @load_store_check {
%z: index) kernel {
%c0 = arith.constant 0 : index
// CHECK-DAG: %[[C4I32:.+]] = arith.constant 4 : i32
+ // CHECK-DAG: %[[C1I32:.+]] = arith.constant 1 : i32
// In-plane offset is 0, so canonicalize folds the add away.
+ // base_height is the row extent size[1] + (size[0] - 1) * batch_rows.
// CHECK: %{{.+}}, %{{.+}}, %[[LSZ:.+]]:3, %[[LSTR:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
- // CHECK: %[[LH:.+]] = arith.index_cast %[[LSZ]]#1 : index to i32
- // CHECK: %[[LBATCH:.+]] = arith.index_cast %[[LSZ]]#0 : index to i32
- // CHECK: %[[LFLAT_H:.+]] = arith.muli %[[LH]], %[[LBATCH]] : i32
// CHECK: %[[LPITCH:.+]] = arith.index_cast %[[LSTR]]#1 : index to i32
// CHECK: %[[LSTRIDE:.+]] = arith.index_cast %[[LSTR]]#0 : index to i32
// CHECK: %[[LROWS:.+]] = arith.divui %[[LSTRIDE]], %[[LPITCH]] : i32
+ // CHECK: %[[LH:.+]] = arith.index_cast %[[LSZ]]#1 : index to i32
+ // CHECK: %[[LBATCH:.+]] = arith.index_cast %[[LSZ]]#0 : index to i32
+ // CHECK: %[[LBATCHM1:.+]] = arith.subi %[[LBATCH]], %[[C1I32]] : i32
+ // CHECK: %[[LBATCH_ROWS:.+]] = arith.muli %[[LBATCHM1]], %[[LROWS]] : i32
+ // CHECK: %[[LFLAT_H:.+]] = arith.addi %[[LH]], %[[LBATCH_ROWS]] : i32
// CHECK: %[[LBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
// CHECK: %[[LZ:.+]] = arith.index_cast %[[Z]] : index to i32
// CHECK: %[[LY:.+]] = arith.muli %[[LZ]], %[[LROWS]] : i32
@@ -118,12 +122,14 @@ gpu.module @load_store_check {
// Store: same handling.
// CHECK: %{{.+}}, %{{.+}}, %[[SSZ:.+]]:3, %[[SSTR:.+]]:3 = memref.extract_strided_metadata %[[DST]]
- // CHECK: %[[SH:.+]] = arith.index_cast %[[SSZ]]#1 : index to i32
- // CHECK: %[[SBATCH:.+]] = arith.index_cast %[[SSZ]]#0 : index to i32
- // CHECK: %[[SFLAT_H:.+]] = arith.muli %[[SH]], %[[SBATCH]] : i32
// CHECK: %[[SPITCH:.+]] = arith.index_cast %[[SSTR]]#1 : index to i32
// CHECK: %[[SSTRIDE:.+]] = arith.index_cast %[[SSTR]]#0 : index to i32
// CHECK: %[[SROWS:.+]] = arith.divui %[[SSTRIDE]], %[[SPITCH]] : i32
+ // CHECK: %[[SH:.+]] = arith.index_cast %[[SSZ]]#1 : index to i32
+ // CHECK: %[[SBATCH:.+]] = arith.index_cast %[[SSZ]]#0 : index to i32
+ // CHECK: %[[SBATCHM1:.+]] = arith.subi %[[SBATCH]], %[[C1I32]] : i32
+ // CHECK: %[[SBATCH_ROWS:.+]] = arith.muli %[[SBATCHM1]], %[[SROWS]] : i32
+ // CHECK: %[[SFLAT_H:.+]] = arith.addi %[[SH]], %[[SBATCH_ROWS]] : i32
// CHECK: %[[SBASE:.+]] = vector.extract %{{.+}}[0] : i64 from vector<4xi64>
// CHECK: %[[SZ:.+]] = arith.index_cast %[[Z]] : index to i32
// CHECK: %[[SY:.+]] = arith.muli %[[SZ]], %[[SROWS]] : i32
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/store_nd_3d_subview.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/store_nd_3d_subview.mlir
new file mode 100644
index 00000000000000..fd1b4878ac67f5
--- /dev/null
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/store_nd_3d_subview.mlir
@@ -0,0 +1,99 @@
+// RUN: mlir-opt %s --gpu-lower-to-xevm-pipeline="xegpu-op-level=workgroup" \
+// RUN: | mlir-runner \
+// RUN: --shared-libs=%mlir_levelzero_runtime \
+// RUN: --shared-libs=%mlir_runner_utils \
+// RUN: --shared-libs=%mlir_c_runner_utils \
+// RUN: --entry-point-result=void \
+// RUN: | FileCheck %s
+
+// A rank-3 descriptor over a subview of a wider source. The 4 planes of the
+// 4x32x32 subview sit 64 rows apart in the 4x64x64 source, so they are not
+// packed back to back: the flattened-plane `base_height` the XeGPUToXeVM
+// lowering hands to the 2D block store is the row extent 32 + 3 * 64 = 224, and
+// every plane has to land inside it for the store to survive the HW boundary
+// check.
+
+module @store_nd_3d_subview attributes {gpu.container_module} {
+ gpu.module @kernel {
+ gpu.func @store_tile(%dst: memref<4x64x64xf32>) kernel {
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant dense<1.000000e+00> : vector<4x32x32xf32>
+ %sub = memref.subview %dst[0, 0, 0] [4, 32, 32] [1, 1, 1]
+ : memref<4x64x64xf32> to memref<4x32x32xf32, strided<[4096, 64, 1]>>
+ %tdesc = xegpu.create_nd_tdesc %sub
+ : memref<4x32x32xf32, strided<[4096, 64, 1]>>
+ -> !xegpu.tensor_desc<4x32x32xf32>
+ xegpu.store_nd %cst, %tdesc[%c0, %c0, %c0]
+ <{layout = #xegpu.layout<sg_layout = [1, 1, 1], sg_data = [4, 32, 32]>}>
+ : vector<4x32x32xf32>, !xegpu.tensor_desc<4x32x32xf32>
+ gpu.return
+ }
+ }
+
+ func.func @test(%dst: memref<4x64x64xf32>) {
+ %c1 = arith.constant 1 : index
+ %c16 = arith.constant 16 : index
+ %dst_gpu = gpu.alloc () : memref<4x64x64xf32>
+ gpu.memcpy %dst_gpu, %dst : memref<4x64x64xf32>, memref<4x64x64xf32>
+ gpu.launch_func @kernel::@store_tile
+ blocks in (%c1, %c1, %c1) threads in (%c16, %c1, %c1)
+ args(%dst_gpu : memref<4x64x64xf32>)
+ gpu.wait
+ gpu.memcpy %dst, %dst_gpu : memref<4x64x64xf32>, memref<4x64x64xf32>
+ gpu.dealloc %dst_gpu : memref<4x64x64xf32>
+ return
+ }
+
+ func.func @main() attributes {llvm.emit_c_interface} {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c4 = arith.constant 4 : index
+ %c32 = arith.constant 32 : index
+ %c64 = arith.constant 64 : index
+ %f0 = arith.constant 0.000000e+00 : f32
+ // Everything the store does not reach keeps this sentinel.
+ %sentinel = arith.constant 9.000000e+00 : f32
+ %dst = memref.alloc() : memref<4x64x64xf32>
+ scf.for %c = %c0 to %c4 step %c1 {
+ scf.for %h = %c0 to %c64 step %c1 {
+ scf.for %w = %c0 to %c64 step %c1 {
+ memref.store %sentinel, %dst[%c, %h, %w] : memref<4x64x64xf32>
+ }
+ }
+ }
+ call @test(%dst) : (memref<4x64x64xf32>) -> ()
+
+ // Sum over the stored tile: 4 * 32 * 32 elements of 1.0. A plane the store
+ // did not reach shows up as 32 * 32 * 9.0 too much.
+ %tile_sum = scf.for %c = %c0 to %c4 step %c1 iter_args(%acc = %f0) -> f32 {
+ %r = scf.for %h = %c0 to %c32 step %c1 iter_args(%acc_h = %acc) -> f32 {
+ %r_w = scf.for %w = %c0 to %c32 step %c1 iter_args(%acc_w = %acc_h) -> f32 {
+ %v = memref.load %dst[%c, %h, %w] : memref<4x64x64xf32>
+ %s = arith.addf %acc_w, %v : f32
+ scf.yield %s : f32
+ }
+ scf.yield %r_w : f32
+ }
+ scf.yield %r : f32
+ }
+ // Sum over the whole buffer, so a store that spilled outside the tile is
+ // caught too: 4096 * 1.0 + 12288 * 9.0.
+ %all_sum = scf.for %c = %c0 to %c4 step %c1 iter_args(%acc = %f0) -> f32 {
+ %r = scf.for %h = %c0 to %c64 step %c1 iter_args(%acc_h = %acc) -> f32 {
+ %r_w = scf.for %w = %c0 to %c64 step %c1 iter_args(%acc_w = %acc_h) -> f32 {
+ %v = memref.load %dst[%c, %h, %w] : memref<4x64x64xf32>
+ %s = arith.addf %acc_w, %v : f32
+ scf.yield %s : f32
+ }
+ scf.yield %r_w : f32
+ }
+ scf.yield %r : f32
+ }
+ // CHECK: 4096
+ vector.print %tile_sum : f32
+ // CHECK: 114688
+ vector.print %all_sum : f32
+ memref.dealloc %dst : memref<4x64x64xf32>
+ return
+ }
+}
More information about the Mlir-commits
mailing list