[Mlir-commits] [mlir] 1e9c8ce - [mlir][xegpu] Lower 1D vector transfers to scattered load/store (#213469)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 3 21:59:57 PDT 2026
Author: Jianhui Li
Date: 2026-08-03T21:59:53-07:00
New Revision: 1e9c8ce9c6aca205a7cff2e3530ab759a1b691f1
URL: https://github.com/llvm/llvm-project/commit/1e9c8ce9c6aca205a7cff2e3530ab759a1b691f1
DIFF: https://github.com/llvm/llvm-project/commit/1e9c8ce9c6aca205a7cff2e3530ab759a1b691f1.diff
LOG: [mlir][xegpu] Lower 1D vector transfers to scattered load/store (#213469)
Both 1D vector.transfer_read / vector.transfer_write now lower to the
scattered xegpu.load / xegpu.store path instead of xegpu.load_nd /
xegpu.store_nd. We reserve the nd block path for rank ≥ 2 vectors.
Cleanup Test: Where both RUN configurations (with and without
--xevm-attach-target) produce identical IR for a function, the
duplicated LOAD-ND/LOAD-GATHER (and STORE-ND/STORE-SCATTER) check blocks
are collapsed into a single shared CHECK block.
---------
Co-authored-by: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 727c98aef6619..795a9e03aa4f8 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -615,12 +615,13 @@ struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
AffineMap readMap = readOp.getPermutationMap();
bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
- // Prefer an nd block load. It requires HW block-load support, a non-0D
- // vector backed by a scalar-element memref, and a map the block load can
- // realize. Out-of-bounds reads are allowed as long as the padding matches
- // load_nd's implicit zero padding.
+ // Prefer an nd block load. It requires HW block-load support, a vector of
+ // rank >= 2 backed by a scalar-element memref, and a map the block load can
+ // realize. 1D vectors use the scattered xegpu.load path instead, which has
+ // a richer interface (e.g. layout capabilities). Out-of-bounds reads are
+ // allowed as long as the padding matches load_nd's implicit zero padding.
bool canLowerToLoadNd =
- hasBlockLoadSupport && loadedVecTy.getRank() > 0 &&
+ hasBlockLoadSupport && loadedVecTy.getRank() > 1 &&
(readMap.isMinorIdentity() || isTransposeLoad) &&
readMemTy.getElementType().isIntOrFloat() &&
(!isOutOfBounds || isZeroOrPoisonPadding(readOp.getPadding()));
@@ -726,12 +727,13 @@ struct TransferWriteLowering
bool hasBlockStoreSupport =
(chip == "pvc" || chip == "bmg" || chip == "cri");
- // Prefer an nd block store. It requires HW block-store support, a non-0D
- // vector backed by a scalar-element memref, and a minor-identity map (block
- // stores have no transpose support). Out-of-bounds writes are handled by
- // the descriptor's boundary check.
+ // Prefer an nd block store. It requires HW block-store support, a vector of
+ // rank >= 2 backed by a scalar-element memref, and a minor-identity map
+ // (block stores have no transpose support). 1D vectors use the scattered
+ // xegpu.store path instead, which has a richer interface. Out-of-bounds
+ // writes are handled by the descriptor's boundary check.
AffineMap map = writeOp.getPermutationMap();
- bool canLowerToStoreNd = hasBlockStoreSupport && vecTy.getRank() > 0 &&
+ bool canLowerToStoreNd = hasBlockStoreSupport && vecTy.getRank() > 1 &&
map.isMinorIdentity() &&
writeMemTy.getElementType().isIntOrFloat();
diff --git a/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
index 2a319869a7b06..9617d347af2b6 100644
--- a/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
@@ -15,8 +15,8 @@ gpu.func @load_1D_vector(%source: memref<8x16x32xf32>,
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK : arith.muli {{.*}} : index
+// CHECK : arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
@@ -42,8 +42,8 @@ gpu.func @load_2D_memref(%source: memref<8x32xf32>,
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {
-// CHECK-COUNT1: arith.muli {{.*}} : index
-// CHECK-COUNT1: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x32xf32> -> index
@@ -69,8 +69,8 @@ gpu.func @load_2D_vector(%source: memref<8x16x32xf32>,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
@@ -97,8 +97,8 @@ gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {
// CHECK: memref.extract_strided_metadata %[[SRC]]
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?xf32> -> index
@@ -125,8 +125,8 @@ gpu.func @load_dynamic_source2(%source: memref<?x8x16xf32>,
// CHECK-SAME: %[[MASK:.+]]: vector<8x16xi1>
// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8x16xf32>) -> vector<8x16xf32> {
// CHECK-NOT: memref.extract_strided_metadata %[[SRC]]
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x8x16xf32> -> index
@@ -238,8 +238,8 @@ gpu.func @non_unit_inner_stride_3D(
// CHECK: %[[BB:.+]], %[[M_OFF:.+]], %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
// CHECK: arith.muli %[[OFF0]], %[[STRIDES]]#0 : index
// CHECK: arith.addi {{.*}} : index
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[STRD_INDICES:.+]] = arith.muli {{.*}}%[[INDICES]]{{.*}} : vector<8xindex>
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[STRD_INDICES]] : vector<8xindex>
diff --git a/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
index ffd3f170c0fad..b9dd241990222 100644
--- a/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
@@ -12,8 +12,8 @@ gpu.func @store_1D_vector(%vec: vector<8xf32>, %source: memref<8x16x32xf32>,
// CHECK-SAME: %[[VAL:.+]]: vector<8xf32>, %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
@@ -35,8 +35,8 @@ gpu.func @store_2D_memref(%vec: vector<8xf32>, %source: memref<8x32xf32>,
// CHECK-SAME: %[[VAL:.+]]: vector<8xf32>, %[[SRC:.+]]: memref<8x32xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
// CHECK-SAME: %[[INDICES:.+]]: vector<8xindex>, %[[MASK:.+]]: vector<8xi1>) {
-// CHECK-COUNT1: arith.muli {{.*}} : index
-// CHECK-COUNT1: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8xindex>
// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x32xf32> -> index
@@ -58,8 +58,8 @@ gpu.func @store_2D_vector(%vec: vector<8x16xf32>, %source: memref<8x16x32xf32>,
// CHECK-SAME: %[[VAL:.+]]: vector<8x16xf32>, %[[SRC:.+]]: memref<8x16x32xf32>,
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>, %[[MASK:.+]]: vector<8x16xi1>) {
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
@@ -82,8 +82,8 @@ gpu.func @store_dynamic_source(%vec: vector<8x16xf32>, %source: memref<?x?x?xf32
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>, %[[MASK:.+]]: vector<8x16xi1>) {
// CHECK: memref.extract_strided_metadata %[[SRC]]
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x?x?xf32> -> index
@@ -106,8 +106,8 @@ gpu.func @store_dynamic_source2(%vec: vector<8x16xf32>, %source: memref<?x8x16xf
// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
// CHECK-SAME: %[[INDICES:.+]]: vector<8x16xindex>, %[[MASK:.+]]: vector<8x16xi1>) {
// CHECK-NOT: memref.extract_strided_metadata %[[SRC]]
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8x16xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[INDICES]] : vector<8x16xindex>
// CHECK: %[[BASE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<?x8x16xf32> -> index
@@ -161,8 +161,8 @@ gpu.func @non_unit_inner_stride_3D(
// CHECK: %[[BB:.+]], %[[M_OFF:.+]], %[[SIZES:.+]]:3, %[[STRIDES:.+]]:3 = memref.extract_strided_metadata %[[SRC]]
// CHECK: arith.muli %[[OFF0]], %[[STRIDES]]#0 : index
// CHECK: arith.addi {{.*}} : index
-// CHECK-COUNT2: arith.muli {{.*}} : index
-// CHECK-COUNT2: arith.addi {{.*}} : index
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi {{.*}} : index
// CHECK: %[[STRD_INDICES:.+]] = arith.muli {{.*}}%[[INDICES]]{{.*}} : vector<8xindex>
// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[STRD_INDICES]] : vector<8xindex>
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index e22e682053a2a..4cc4a0db5b63c 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -9,36 +9,17 @@ gpu.func @load_1D_vector(%source: memref<8x16x32xf32>, %offset: index) -> vector
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @load_1D_vector(
-// LOAD-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// LOAD-ND-SAME: %[[OFFSET:.+]]: index
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFF1:.*]], %[[SIZES:.*]], %[[STRIDES:.*]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// LOAD-ND-SAME: : memref<32xf32, strided<[1], offset: ?>> -> memref<f32>, index, index, index
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// LOAD-ND-SAME: : memref<f32> -> index
-// LOAD-ND: %[[MUL:.*]] = arith.muli %[[OFF1]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.*]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [32],
-// LOAD-ND-SAME: strides : [1] : i64 -> !xegpu.tensor_desc<8xf32,
-// LOAD-ND-SAME: #xegpu.block_tdesc_attr<boundary_check = false>>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFFSET]]]
-// LOAD-ND-SAME: : !xegpu.tensor_desc<8xf32, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<8xf32>
-
-// LOAD-GATHER-LABEL: @load_1D_vector(
-// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// LOAD-GATHER: %[[STEP:.+]] = vector.step : vector<8xindex>
-// LOAD-GATHER-COUNT2: arith.muli {{.*}} : index
-// LOAD-GATHER-COUNT2: arith.addi {{.*}} : index
-// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
-// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
-// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
-// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
-
+// CHECK-LABEL: @load_1D_vector(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
+// CHECK-COUNT-2: arith.muli {{.*}} : index
+// CHECK-COUNT-2: arith.addi {{.*}} : index
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
}
// -----
@@ -81,10 +62,8 @@ gpu.func @load_2D_vector(%source: memref<8x16x32xf32>,
// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
-
}
-
// -----
gpu.module @xevm_module {
gpu.func @load_zero_pad_out_of_bounds(%source: memref<32x64xf32>,
@@ -105,10 +84,8 @@ gpu.func @load_zero_pad_out_of_bounds(%source: memref<32x64xf32>,
// LOAD-GATHER-LABEL: @load_zero_pad_out_of_bounds(
// LOAD-GATHER: vector.transfer_read
-
}
-
// -----
gpu.module @xevm_module {
gpu.func @load_transposed(%source: memref<32x64xf32>,
@@ -131,7 +108,6 @@ gpu.func @load_transposed(%source: memref<32x64xf32>,
// LOAD-ND: %[[VEC_TRANSPOSED:.+]] = vector.transpose %[[VEC]], [1, 0] : vector<16x8xf32> to vector<8x16xf32>
// LOAD-ND: return %[[VEC_TRANSPOSED]]
-
// LOAD-GATHER-LABEL: @load_transposed(
// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<32x64xf32>,
// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
@@ -145,7 +121,6 @@ gpu.func @load_transposed(%source: memref<32x64xf32>,
// LOAD-GATHER: %[[COLLAPSE:.*]] = memref.extract_aligned_pointer_as_index %arg0 : memref<32x64xf32> -> index
// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// LOAD-GATHER: %[[LOAD:.*]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
-
}
// -----
@@ -185,7 +160,6 @@ gpu.func @load_transpose_3d_memref(%source: memref<32x64x128xf32>,
// LOAD-GATHER: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32x64x128xf32> -> index
// LOAD-GATHER-NEXT: %[[I64PTR:.+]] = arith.index_cast %[[INTPTR]] : index to i64
// LOAD-GATHER-NEXT: %[[LOAD:.*]] = xegpu.load %[[I64PTR]][%[[IDX]]], %{{.*}} : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
-
}
// -----
@@ -240,7 +214,6 @@ gpu.func @load_dynamic_source(%source: memref<?x?x?xf32>,
// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF1]], %[[OFF2]]]{{.*}}-> vector<8x16xf32>
// LOAD-ND: return %[[VEC]]
-
// LOAD-GATHER-LABEL: @load_dynamic_source(
// LOAD-GATHER-SAME: %[[ARG0:.+]]: memref<?x?x?xf32>,
// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
@@ -295,7 +268,6 @@ gpu.func @load_dynamic_source2(%source: memref<?x8x16xf32>,
// LOAD-GATHER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %arg0 : memref<?x8x16xf32> -> index
// LOAD-GATHER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[OFFSETS]]{{\]}}, %[[CST_0]] : i64, vector<8x16xindex>, vector<8x16xi1> -> vector<8x16xf32>
-
}
// -----
@@ -366,7 +338,6 @@ gpu.func @load_high_dim_vector(%source: memref<16x32x64xf32>,
// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %arg0 : memref<16x32x64xf32> -> index
// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<8x16x32xindex>, vector<8x16x32xi1> -> vector<8x16x32xf32>
-
}
// -----
@@ -398,7 +369,6 @@ gpu.func @load_8D_vector(%source: memref<2x2x2x2x2x2x2x2xf32>,
// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<2x2x2x2x2x2x2x2xf32> -> index
// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]][%[[IDX]]], %[[CST]] : i64, vector<2x2x2x2x2x2x2x2xindex>, vector<2x2x2x2x2x2x2x2xi1> -> vector<2x2x2x2x2x2x2x2xf32>
-
}
// -----
@@ -443,11 +413,8 @@ gpu.func @no_load_out_of_bounds_non_zero_pad(%source: memref<32x64xf32>,
gpu.return %0, %1 : vector<8x16xf32>, vector<8x16xf32>
}
-// LOAD-ND-LABEL: @no_load_out_of_bounds_non_zero_pad(
-// LOAD-ND-COUNT-2: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @no_load_out_of_bounds_non_zero_pad(
-// LOAD-GATHER-COUNT-2: vector.transfer_read
+// CHECK-LABEL: @no_load_out_of_bounds_non_zero_pad(
+// CHECK-COUNT-2: vector.transfer_read
}
// -----
@@ -460,11 +427,8 @@ gpu.func @no_load_out_of_bounds_1D_vector(%source: memref<8x16x32xf32>,
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @no_load_out_of_bounds_1D_vector(
-// LOAD-ND: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @no_load_out_of_bounds_1D_vector(
-// LOAD-GATHER: vector.transfer_read
+// CHECK-LABEL: @no_load_out_of_bounds_1D_vector(
+// CHECK: vector.transfer_read
}
// -----
@@ -478,11 +442,8 @@ gpu.func @no_load_masked(%source : memref<4xf32>,
gpu.return %0 : vector<4xf32>
}
-// LOAD-ND-LABEL: @no_load_masked(
-// LOAD-ND: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @no_load_masked(
-// LOAD-GATHER: vector.transfer_read
+// CHECK-LABEL: @no_load_masked(
+// CHECK: vector.transfer_read
}
// -----
@@ -495,14 +456,10 @@ gpu.func @no_load_tensor(%source: tensor<32x64xf32>,
gpu.return %0 : vector<8x16xf32>
}
-// LOAD-ND-LABEL: @no_load_tensor(
-// LOAD-ND: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @no_load_tensor(
-// LOAD-GATHER: vector.transfer_read
+// CHECK-LABEL: @no_load_tensor(
+// CHECK: vector.transfer_read
}
-
// -----
gpu.module @xevm_module {
gpu.func @no_load_non_unit_inner_stride(
@@ -514,14 +471,10 @@ gpu.func @no_load_non_unit_inner_stride(
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @no_load_non_unit_inner_stride(
-// LOAD-ND: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @no_load_non_unit_inner_stride(
-// LOAD-GATHER: vector.transfer_read
+// CHECK-LABEL: @no_load_non_unit_inner_stride(
+// CHECK: vector.transfer_read
}
-
// -----
gpu.module @xevm_module {
gpu.func @no_load_unsupported_map(%source: memref<16x32x64xf32>,
@@ -533,11 +486,8 @@ gpu.func @no_load_unsupported_map(%source: memref<16x32x64xf32>,
gpu.return %0 : vector<8x16xf32>
}
-// LOAD-ND-LABEL: @no_load_unsupported_map(
-// LOAD-ND: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @no_load_unsupported_map(
-// LOAD-GATHER: vector.transfer_read
+// CHECK-LABEL: @no_load_unsupported_map(
+// CHECK: vector.transfer_read
}
// -----
@@ -550,36 +500,21 @@ gpu.func @load_from_subview_1D(%source: memref<4096x4096xf16>, %off1: index, %of
gpu.return %0 : vector<8xf16>
}
-// LOAD-ND-LABEL: @load_from_subview_1D(
-// LOAD-ND-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
-// LOAD-ND-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// LOAD-ND: %[[ELEM_BYTES:.+]] = arith.constant 2 : index
-// LOAD-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// LOAD-ND: %[[COLLAPSED:.+]] = memref.subview %[[SUBVIEW]][%[[OFF2]], 0]
-// LOAD-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.*]], %[[STRIDES:.*]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// LOAD-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// LOAD-ND: %[[MUL:.+]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// LOAD-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// LOAD-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// LOAD-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [256], strides : [1] : i64 ->
-// LOAD-ND-SAME: !xegpu.tensor_desc<8xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
-// LOAD-ND: %[[VEC:.+]] = xegpu.load_nd %[[DESC]][%[[OFF2]]] : !xegpu.tensor_desc<8xf16, #xegpu.block_tdesc_attr<boundary_check = false>> -> vector<8xf16>
-
-// LOAD-GATHER-LABEL: @load_from_subview_1D(
-// LOAD-GATHER-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
-// LOAD-GATHER-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// LOAD-GATHER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// LOAD-GATHER: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// LOAD-GATHER: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
-// LOAD-GATHER: %[[STEP:.+]] = vector.step : vector<8xindex>
-// LOAD-GATHER: arith.muli {{.*}} : index
-// LOAD-GATHER: arith.addi %[[OFFSET]]{{.*}} : index
-// LOAD-GATHER: arith.addi {{.*}} : index
-// LOAD-GATHER: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
-// LOAD-GATHER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
-// LOAD-GATHER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
-// LOAD-GATHER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// LOAD-GATHER: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf16>
+// CHECK-LABEL: @load_from_subview_1D(
+// CHECK-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
+// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
+// CHECK: %[[BB:.+]], %[[OFFSET:.+]],{{.*}},{{.*}} = memref.extract_strided_metadata %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
+// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi %[[OFFSET]]{{.*}} : index
+// CHECK: arith.addi {{.*}} : index
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]] : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf16>
}
// -----
@@ -635,20 +570,12 @@ gpu.func @load_2D_vector_addrspace3(%source: memref<16x32xf32, 3>,
gpu.return %0 : vector<8x16xf32>
}
-// LOAD-ND-LABEL: @load_2D_vector_addrspace3
-// LOAD-ND-SAME: %[[SOURCE:.+]]: memref<16x32xf32, 3>
-// LOAD-ND-SAME: %[[OFFSET:.+]]: index
-// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
-// LOAD-ND: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]], %[[OFFSET]]] : !xegpu.mem_desc<16x32xf32>, index, index -> vector<8x16xf32>
-// LOAD-ND: gpu.return %[[DATA]] : vector<8x16xf32>
-
-// LOAD-GATHER-LABEL: @load_2D_vector_addrspace3
-// LOAD-GATHER-SAME: %[[SOURCE:.+]]: memref<16x32xf32, 3>
-// LOAD-GATHER-SAME: %[[OFFSET:.+]]: index
-// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
-// LOAD-GATHER: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]], %[[OFFSET]]] : !xegpu.mem_desc<16x32xf32>, index, index -> vector<8x16xf32>
-// LOAD-GATHER: gpu.return %[[DATA]] : vector<8x16xf32>
-
+// CHECK-LABEL: @load_2D_vector_addrspace3
+// CHECK-SAME: %[[SOURCE:.+]]: memref<16x32xf32, 3>
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
+// CHECK: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]], %[[OFFSET]]] : !xegpu.mem_desc<16x32xf32>, index, index -> vector<8x16xf32>
+// CHECK: gpu.return %[[DATA]] : vector<8x16xf32>
}
// -----
@@ -661,20 +588,12 @@ gpu.func @load_1D_vector_addrspace3(%source: memref<32xf32, 3>,
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @load_1D_vector_addrspace3
-// LOAD-ND-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
-// LOAD-ND-SAME: %[[OFFSET:.+]]: index
-// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
-// LOAD-ND: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]]] : !xegpu.mem_desc<32xf32>, index -> vector<8xf32>
-// LOAD-ND: gpu.return %[[DATA]] : vector<8xf32>
-
-// LOAD-GATHER-LABEL: @load_1D_vector_addrspace3
-// LOAD-GATHER-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
-// LOAD-GATHER-SAME: %[[OFFSET:.+]]: index
-// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
-// LOAD-GATHER: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]]] : !xegpu.mem_desc<32xf32>, index -> vector<8xf32>
-// LOAD-GATHER: gpu.return %[[DATA]] : vector<8xf32>
-
+// CHECK-LABEL: @load_1D_vector_addrspace3
+// CHECK-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
+// CHECK: %[[DATA:.+]] = xegpu.load_matrix %[[MEM_DESC]][%[[OFFSET]]] : !xegpu.mem_desc<32xf32>, index -> vector<8xf32>
+// CHECK: gpu.return %[[DATA]] : vector<8xf32>
}
// -----
@@ -691,16 +610,10 @@ gpu.func @load_2D_vector_alloca_promoted_to_slm(%offset: index)
gpu.return %0 : vector<8x16xf32>
}
-// LOAD-ND-LABEL: @load_2D_vector_alloca_promoted_to_slm
-// LOAD-ND: %[[BUF:.+]] = memref.alloca() : memref<16x32xf32, 3>
-// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
-// LOAD-ND: xegpu.load_matrix %[[MEM_DESC]]
-
-// LOAD-GATHER-LABEL: @load_2D_vector_alloca_promoted_to_slm
-// LOAD-GATHER: %[[BUF:.+]] = memref.alloca() : memref<16x32xf32, 3>
-// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
-// LOAD-GATHER: xegpu.load_matrix %[[MEM_DESC]]
-
+// CHECK-LABEL: @load_2D_vector_alloca_promoted_to_slm
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<16x32xf32, 3>
+// CHECK: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
+// CHECK: xegpu.load_matrix %[[MEM_DESC]]
}
// -----
@@ -716,16 +629,10 @@ gpu.func @load_1D_vector_alloca_promoted_to_slm(%offset: index)
gpu.return %0 : vector<8xf32>
}
-// LOAD-ND-LABEL: @load_1D_vector_alloca_promoted_to_slm
-// LOAD-ND: %[[BUF:.+]] = memref.alloca() : memref<16xf32, 3>
-// LOAD-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16xf32, 3> -> !xegpu.mem_desc<16xf32>
-// LOAD-ND: xegpu.load_matrix %[[MEM_DESC]]
-
-// LOAD-GATHER-LABEL: @load_1D_vector_alloca_promoted_to_slm
-// LOAD-GATHER: %[[BUF:.+]] = memref.alloca() : memref<16xf32, 3>
-// LOAD-GATHER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16xf32, 3> -> !xegpu.mem_desc<16xf32>
-// LOAD-GATHER: xegpu.load_matrix %[[MEM_DESC]]
-
+// CHECK-LABEL: @load_1D_vector_alloca_promoted_to_slm
+// CHECK: %[[BUF:.+]] = memref.alloca() : memref<16xf32, 3>
+// CHECK: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[BUF]] : memref<16xf32, 3> -> !xegpu.mem_desc<16xf32>
+// CHECK: xegpu.load_matrix %[[MEM_DESC]]
}
// -----
@@ -738,7 +645,6 @@ gpu.func @load_0D_memref_unsupported(%source: memref<f16>) -> vector<f16> {
// CHECK-LABEL: @load_0D_memref_unsupported
// CHECK: vector.transfer_read
-
}
// -----
@@ -751,12 +657,8 @@ gpu.func @load_0D_vector_unsupported(%source: memref<3xf32>,
gpu.return %0 : vector<f32>
}
-// LOAD-ND-LABEL: @load_0D_vector_unsupported
-// LOAD-ND: vector.transfer_read
-
-// LOAD-GATHER-LABEL: @load_0D_vector_unsupported
-// LOAD-GATHER: vector.transfer_read
-
+// CHECK-LABEL: @load_0D_vector_unsupported
+// CHECK: vector.transfer_read
}
// -----
@@ -816,5 +718,4 @@ gpu.func @transpose_1x1024x24x64(
// LOAD-GATHER: arith.muli %block_id_x, %[[C65536]] : index
// LOAD-GATHER: arith.muli %{{.+}}, %[[C64]] : index
// LOAD-GATHER: xegpu.store {{.*}}
-
}
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index 2d55a785595f7..dad2740f5c0ea 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -1,7 +1,6 @@
// RUN: mlir-opt %s --xevm-attach-target='module=xevm_* O=3 chip=pvc' -convert-vector-to-xegpu -split-input-file | FileCheck %s --check-prefixes=STORE-ND,CHECK
// RUN: mlir-opt %s -convert-vector-to-xegpu -split-input-file | FileCheck %s --check-prefixes=STORE-SCATTER,CHECK
-
gpu.module @xevm_module {
gpu.func @store_1D_vector(%vec: vector<8xf32>,
%source: memref<8x16x32xf32>, %offset: index) {
@@ -11,36 +10,18 @@ gpu.func @store_1D_vector(%vec: vector<8xf32>,
gpu.return
}
-// STORE-ND-LABEL: @store_1D_vector(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[ELEM_BYTES:.+]] = arith.constant 4 : index
-// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SRC]][%[[OFFSET]], %[[OFFSET]], 0]
-// STORE-ND: %[[BASE_BUFFER:.+]], %[[OFFSET1:.+]], %[[SIZES:.+]], %[[STRIDES:.+]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// STORE-ND-SAME: : memref<32xf32, strided<[1], offset: ?>> -> memref<f32>, index, index, index
-// STORE-ND: %[[INTPTR:.+]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// STORE-ND-SAME: : memref<f32> -> index
-// STORE-ND: %[[MUL:.+]] = arith.muli %[[OFFSET1]], %[[ELEM_BYTES]] : index
-// STORE-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// STORE-ND: %[[I64PTR:.+]] = arith.index_cast %[[ADD]] : index to i64
-// STORE-ND: %[[DESC:.+]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [32],
-// STORE-ND-SAME: strides : [1] : i64 -> !xegpu.tensor_desc<8xf32,
-// STORE-ND-SAME: boundary_check = false
-// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFFSET]]] : vector<8xf32>
-
-// STORE-SCATTER-LABEL: @store_1D_vector(
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf32>,
-// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
-// STORE-SCATTER-DAG: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// STORE-SCATTER-DAG: %[[STEP:.+]] = vector.step
-// STORE-SCATTER-COUNT2: arith.muli {{.*}} : index
-// STORE-SCATTER-COUNT2: arith.addi {{.*}} : index
-// STORE-SCATTER-DAG: %[[BCAST:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
-// STORE-SCATTER-DAG: %[[IDX:.+]] = arith.addi %[[BCAST]], %{{.*}} : vector<8xindex>
-// STORE-SCATTER-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
-// STORE-SCATTER-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>
+// CHECK-LABEL: @store_1D_vector(
+// CHECK-SAME: %[[VEC:.+]]: vector<8xf32>,
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-DAG: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// CHECK-DAG: %[[STEP:.+]] = vector.step
+// CHECK : arith.muli {{.*}} : index
+// CHECK : arith.addi {{.*}} : index
+// CHECK-DAG: %[[BCAST:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
+// CHECK-DAG: %[[IDX:.+]] = arith.addi %[[BCAST]], %{{.*}} : vector<8xindex>
+// CHECK-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<8x16x32xf32> -> index
+// CHECK-DAG: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>
}
// -----
@@ -166,9 +147,9 @@ gpu.func @store_transposed(%vec: vector<8x16xf32>,
// CHECK-SAME: %[[SRC:.+]]: memref<32x64xf32>,
// CHECK-SAME: %[[OFFSET:.+]]: index
// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8x16xi1>
-// CHECK-COUNT2: %[[STEP:.+]] = vector.step
-// CHECK-COUNT2: vector.shape_cast {{.*}}
-// CHECK-COUNT2: vector.broadcast {{.*}} : vector<8x16xindex>
+// CHECK : %[[STEP:.+]] = vector.step
+// CHECK : vector.shape_cast {{.*}}
+// CHECK : vector.broadcast {{.*}} : vector<8x16xindex>
// CHECK-DAG: %[[BCAST2:.+]] = vector.broadcast {{.*}} : index to vector<8x16xindex>
// CHECK-DAG: %[[IDX:.+]] = arith.addi %[[BCAST2]], {{.*}} : vector<8x16xindex>
// CHECK-DAG: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SRC]] : memref<32x64xf32> -> index
@@ -274,11 +255,8 @@ gpu.func @no_store_masked(%vec: vector<4xf32>,
gpu.return
}
-// STORE-ND-LABEL: @no_store_masked(
-// STORE-ND: vector.transfer_write
-
-// STORE-SCATTER-LABEL: @no_store_masked(
-// STORE-SCATTER: vector.transfer_write
+// CHECK-LABEL: @no_store_masked(
+// CHECK: vector.transfer_write
}
// -----
@@ -291,11 +269,8 @@ gpu.func @no_store_tensor(%vec: vector<8x16xf32>,
gpu.return %0 : tensor<32x64xf32>
}
-// STORE-ND-LABEL: @no_store_tensor(
-// STORE-ND: vector.transfer_write
-
-// STORE-SCATTER-LABEL: @no_store_tensor(
-// STORE-SCATTER: vector.transfer_write
+// CHECK-LABEL: @no_store_tensor(
+// CHECK: vector.transfer_write
}
// -----
@@ -308,11 +283,8 @@ gpu.func @no_store_non_unit_inner_stride(%vec: vector<8xf32>,
gpu.return
}
-// STORE-ND-LABEL: @no_store_non_unit_inner_stride(
-// STORE-ND: vector.transfer_write
-
-// STORE-SCATTER-LABEL: @no_store_non_unit_inner_stride(
-// STORE-SCATTER: vector.transfer_write
+// CHECK-LABEL: @no_store_non_unit_inner_stride(
+// CHECK: vector.transfer_write
}
// -----
@@ -326,11 +298,8 @@ gpu.func @no_store_unsupported_map(%vec: vector<8x16xf32>,
gpu.return
}
-// STORE-ND-LABEL: @no_store_unsupported_map(
-// STORE-ND: vector.transfer_write
-
-// STORE-SCATTER-LABEL: @no_store_unsupported_map(
-// STORE-SCATTER: vector.transfer_write
+// CHECK-LABEL: @no_store_unsupported_map(
+// CHECK: vector.transfer_write
}
// -----
@@ -343,11 +312,8 @@ gpu.func @no_store_out_of_bounds_1D_vector(%vec: vector<8xf32>,
gpu.return
}
-// STORE-ND-LABEL: @no_store_out_of_bounds_1D_vector(
-// STORE-ND: vector.transfer_write
-
-// STORE-SCATTER-LABEL: @no_store_out_of_bounds_1D_vector(
-// STORE-SCATTER: vector.transfer_write
+// CHECK-LABEL: @no_store_out_of_bounds_1D_vector(
+// CHECK: vector.transfer_write
}
// -----
@@ -362,41 +328,26 @@ gpu.func @store_to_subview(%vec: vector<8xf16>,
: vector<8xf16>, memref<256x256xf16, strided<[4096, 1], offset: ?>>
gpu.return
}
-// STORE-ND-LABEL: @store_to_subview(
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf16>,
-// STORE-ND-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
-// STORE-ND-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// STORE-ND: %[[ELEM_BYTES:.+]] = arith.constant 2 : index
-// STORE-ND: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1] : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// STORE-ND: %[[COLLAPSED:.+]] = memref.subview %[[SUBVIEW]][%[[OFF2]], 0]
-// STORE-ND: %[[BASE_BUFFER:.*]], %[[OFFSET:.*]], %[[SIZES:.*]], %[[STRIDES:.*]] = memref.extract_strided_metadata %[[COLLAPSED]]
-// STORE-ND: %[[INTPTR:.*]] = memref.extract_aligned_pointer_as_index %[[BASE_BUFFER]]
-// STORE-ND: %[[MUL:.+]] = arith.muli %[[OFFSET]], %[[ELEM_BYTES]] : index
-// STORE-ND: %[[ADD:.+]] = arith.addi %[[INTPTR]], %[[MUL]] : index
-// STORE-ND: %[[I64PTR:.*]] = arith.index_cast %[[ADD]] : index to i64
-// STORE-ND: %[[DESC:.*]] = xegpu.create_nd_tdesc %[[I64PTR]], shape : [256], strides : [1] : i64 ->
-// STORE-ND-SAME: !xegpu.tensor_desc<8xf16, #xegpu.block_tdesc_attr<boundary_check = false>>
-// STORE-ND: xegpu.store_nd %[[VEC]], %[[DESC]][%[[OFF2]]] : vector<8xf16>
-
-// STORE-SCATTER-LABEL: @store_to_subview(
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf16>,
-// STORE-SCATTER-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
-// STORE-SCATTER-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
-// STORE-SCATTER: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
-// STORE-SCATTER: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]
-// STORE-SCATTER-SAME: : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
-// STORE-SCATTER: %[[BB:.+]], %[[OFFSET:.+]], {{.*}}, {{.*}} = memref.extract_strided_metadata %[[SUBVIEW]]
-// STORE-SCATTER-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
-// STORE-SCATTER: %[[STEP:.+]] = vector.step : vector<8xindex>
-// STORE-SCATTER: arith.muli {{.*}} : index
-// STORE-SCATTER: arith.addi %[[OFFSET]]{{.*}} : index
-// STORE-SCATTER: arith.addi {{.*}} : index
-// STORE-SCATTER: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
-// STORE-SCATTER: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
-// STORE-SCATTER: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]]
-// STORE-SCATTER-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
-// STORE-SCATTER: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
-// STORE-SCATTER: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>
+
+// CHECK-LABEL: @store_to_subview(
+// CHECK-SAME: %[[VEC:.+]]: vector<8xf16>,
+// CHECK-SAME: %[[SRC:.+]]: memref<4096x4096xf16>,
+// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index
+// CHECK: %[[CST:.+]] = arith.constant dense<true> : vector<8xi1>
+// CHECK: %[[SUBVIEW:.+]] = memref.subview %[[SRC]][%[[OFF1]], %[[OFF2]]] [256, 256] [1, 1]
+// CHECK-SAME: : memref<4096x4096xf16> to memref<256x256xf16, strided<[4096, 1], offset: ?>>
+// CHECK: %[[BB:.+]], %[[OFFSET:.+]], {{.*}}, {{.*}} = memref.extract_strided_metadata %[[SUBVIEW]]
+// CHECK-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> memref<f16>, index, index, index, index, index
+// CHECK: %[[STEP:.+]] = vector.step : vector<8xindex>
+// CHECK: arith.muli {{.*}} : index
+// CHECK: arith.addi %[[OFFSET]]{{.*}} : index
+// CHECK: arith.addi {{.*}} : index
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
+// CHECK: %[[IDX:.+]] = arith.addi %[[SPLAT]], %[[STEP]] : vector<8xindex>
+// CHECK: %[[COLLAPSE:.+]] = memref.extract_aligned_pointer_as_index %[[SUBVIEW]]
+// CHECK-SAME: : memref<256x256xf16, strided<[4096, 1], offset: ?>> -> index
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast %[[COLLAPSE]] : index to i64
+// CHECK: xegpu.store %[[VEC]], %[[COLLAPSE_I]]{{\[}}%[[IDX]]{{\]}}, %[[CST]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>
}
// -----
@@ -409,22 +360,13 @@ gpu.func @store_2D_vector_addrspace3(%vec: vector<8x16xf32>,
gpu.return
}
-// STORE-ND-LABEL: @store_2D_vector_addrspace3
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8x16xf32>
-// STORE-ND-SAME: %[[SOURCE:.+]]: memref<16x32xf32, 3>
-// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
-// STORE-ND: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>, !xegpu.mem_desc<16x32xf32>, index, index
-// STORE-ND: gpu.return
-
-// STORE-SCATTER-LABEL: @store_2D_vector_addrspace3
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8x16xf32>
-// STORE-SCATTER-SAME: %[[SOURCE:.+]]: memref<16x32xf32, 3>
-// STORE-SCATTER-SAME: %[[OFFSET:.+]]: index
-// STORE-SCATTER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
-// STORE-SCATTER: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>, !xegpu.mem_desc<16x32xf32>, index, index
-// STORE-SCATTER: gpu.return
-
+// CHECK-LABEL: @store_2D_vector_addrspace3
+// CHECK-SAME: %[[VEC:.+]]: vector<8x16xf32>
+// CHECK-SAME: %[[SOURCE:.+]]: memref<16x32xf32, 3>
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<16x32xf32, 3> -> !xegpu.mem_desc<16x32xf32>
+// CHECK: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]], %[[OFFSET]]] : vector<8x16xf32>, !xegpu.mem_desc<16x32xf32>, index, index
+// CHECK: gpu.return
}
// -----
@@ -437,22 +379,13 @@ gpu.func @store_1D_vector_addrspace3(%vec: vector<8xf32>,
gpu.return
}
-// STORE-ND-LABEL: @store_1D_vector_addrspace3
-// STORE-ND-SAME: %[[VEC:.+]]: vector<8xf32>
-// STORE-ND-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
-// STORE-ND-SAME: %[[OFFSET:.+]]: index
-// STORE-ND: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
-// STORE-ND: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]]] : vector<8xf32>, !xegpu.mem_desc<32xf32>, index
-// STORE-ND: gpu.return
-
-// STORE-SCATTER-LABEL: @store_1D_vector_addrspace3
-// STORE-SCATTER-SAME: %[[VEC:.+]]: vector<8xf32>
-// STORE-SCATTER-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
-// STORE-SCATTER-SAME: %[[OFFSET:.+]]: index
-// STORE-SCATTER: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
-// STORE-SCATTER: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]]] : vector<8xf32>, !xegpu.mem_desc<32xf32>, index
-// STORE-SCATTER: gpu.return
-
+// CHECK-LABEL: @store_1D_vector_addrspace3
+// CHECK-SAME: %[[VEC:.+]]: vector<8xf32>
+// CHECK-SAME: %[[SOURCE:.+]]: memref<32xf32, 3>
+// CHECK-SAME: %[[OFFSET:.+]]: index
+// CHECK: %[[MEM_DESC:.+]] = xegpu.create_mem_desc %[[SOURCE]] : memref<32xf32, 3> -> !xegpu.mem_desc<32xf32>
+// CHECK: xegpu.store_matrix %[[VEC]], %[[MEM_DESC]][%[[OFFSET]]] : vector<8xf32>, !xegpu.mem_desc<32xf32>, index
+// CHECK: gpu.return
}
// -----
@@ -464,10 +397,6 @@ gpu.func @store_0D_vector_unsupported(%vec: vector<f32>,
gpu.return
}
-// STORE-ND-LABEL: @store_0D_vector_unsupported
-// STORE-ND: vector.transfer_write
-
-// STORE-SCATTER-LABEL: @store_0D_vector_unsupported
-// STORE-SCATTER: vector.transfer_write
-
+// CHECK-LABEL: @store_0D_vector_unsupported
+// CHECK: vector.transfer_write
}
More information about the Mlir-commits
mailing list