[Mlir-commits] [mlir] [mlir][xegpu] Lower 1D vector transfers to scattered load/store (PR #213469)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Aug 1 09:38:15 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
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.
---
Patch is 35.92 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/213469.diff
3 Files Affected:
- (modified) mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (+12-10)
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir (+60-159)
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir (+58-129)
``````````diff
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/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index e22e682053a2a..68da1bfeda987 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-COUNT2: arith.muli {{.*}} : index
+// CHECK-COUNT2: 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...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/213469
More information about the Mlir-commits
mailing list