[Mlir-commits] [mlir] [mlir][xegpu] Cast non-index gather/scatter indices before offset math (PR #216492)
Hamza Qureshi
llvmlistbot at llvm.org
Sat Aug 15 07:12:56 PDT 2026
https://github.com/hamzaqureshi5 created https://github.com/llvm/llvm-project/pull/216492
Fixes #205977
`mlir-opt -convert-vector-to-xegpu` crashes on this input:
```mlir
func.func @m(%arg0: memref<4x2xf32>, %arg1: vector<4xi32>, %arg2: vector<4xi1>, %arg3: vector<4xf32>) -> vector<4xf32> {
%c0 = arith.constant 0 : index
%0 = vector.gather %arg0[%c0, %c0] [%arg1], %arg2, %arg3 : memref<4x2xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
return %0 : vector<4xf32>
}
```
```
error: 'arith.addi' op folder produced a value of incorrect type: 'vector<4xi32>', expected: 'vector<4xindex>'
note: see current operation: %8 = "arith.addi"(%6, %7) : (vector<4xi32>, vector<4xindex>) -> vector<4xindex>
Assertion `false && "incorrect fold result type"' failed.
```
To linearize a gather or scatter, `computeOffsets` combines the memref strides and the base offset with the per-element indices. The strides and the base offset are index typed, and the base is broadcast to a `vector<...xindex>`, but the strided indices keep the element type of the op's index vector. `vector.gather` and `vector.scatter` accept any integer type there, so with an `i32` index vector the resulting `arith.addi` has mismatched operand types.
This casts the indices to index type when they use a different integer type, so the offset arithmetic is well typed. The equivalent helper used for transfer read/write already works in index type throughout.
## Testing
All existing tests in `gather-to-xegpu.mlir` and `scatter-to-xegpu.mlir` use `vector<...xindex>` indices, so this case was uncovered. Added an `i32` index case to both, since the two ops share the fixed helper.
Verified on an assertions build that the reproducer aborts before the change and lowers cleanly after, that both new tests fail without the change, that the produced IR round trips through `mlir-opt`, and that `mlir/test/Conversion` stays at 412/412.
>From 5b8a9a49c8fe2bdccd529f8a2b65e057c8d14f96 Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Sat, 15 Aug 2026 19:09:24 +0500
Subject: [PATCH] [mlir][xegpu] Cast non-index gather/scatter indices before
offset math
The strides and the base offset used to linearize a vector.gather or
vector.scatter are computed in index type, but the index vector of those ops
may use any integer type. The lowering combined the two without a conversion,
producing an arith.addi whose operand types disagree, which tripped the fold
result type check.
Cast the indices to index type when they use a different integer type.
Fixes #205977
---
.../VectorToXeGPU/VectorToXeGPU.cpp | 8 ++++++
.../VectorToXeGPU/gather-to-xegpu.mlir | 26 +++++++++++++++++++
.../VectorToXeGPU/scatter-to-xegpu.mlir | 23 ++++++++++++++++
3 files changed, 57 insertions(+)
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 8a45836426931..be9e8e8aa10a2 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -399,6 +399,14 @@ static Value computeOffsets(PatternRewriter &rewriter, OpType gatScatOp,
Value indices = gatScatOp.getIndices();
VectorType vecType = cast<VectorType>(indices.getType());
+ // The strides and the base offset are computed in index type, while the
+ // gather/scatter indices may use any integer type. Cast them to index so
+ // that the offset arithmetic below is well typed.
+ if (!vecType.getElementType().isIndex()) {
+ vecType = VectorType::get(vecType.getShape(), rewriter.getIndexType());
+ indices = arith::IndexCastOp::create(rewriter, loc, vecType, indices);
+ }
+
Value strideVector =
vector::BroadcastOp::create(rewriter, loc, vecType, strides.back())
.getResult();
diff --git a/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
index 9617d347af2b6..16c3c0b34758a 100644
--- a/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
@@ -249,3 +249,29 @@ gpu.func @non_unit_inner_stride_3D(
// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[V]], %[[PASS]] : vector<8xi1>, vector<8xf32>
// CHECK: gpu.return %[[RES]] : vector<8xf32>
}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_1D_vector_i32_indices(%source: memref<8x16x32xf32>,
+ %off1: index, %off2: index, %off3: index,
+ %indices: vector<8xi32>, %mask: vector<8xi1>,
+ %pass_thru: vector<8xf32>) -> vector<8xf32> {
+ %0 = vector.gather %source[%off1, %off2, %off3][%indices], %mask,
+ %pass_thru : memref<8x16x32xf32>, vector<8xi32>, vector<8xi1>, vector<8xf32> into vector<8xf32>
+ gpu.return %0 : vector<8xf32>
+}
+// Indices that are not index typed are cast before the offset arithmetic.
+// CHECK-LABEL: @load_1D_vector_i32_indices(
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
+// CHECK-SAME: %[[INDICES:.+]]: vector<8xi32>
+// CHECK-SAME: %[[MASK:.+]]: vector<8xi1>
+// CHECK-SAME: %[[PASS_THRU:.+]]: vector<8xf32>) -> vector<8xf32> {
+// CHECK: %[[IDX:.+]] = arith.index_cast %[[INDICES]] : vector<8xi32> to vector<8xindex>
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}}: index to vector<8xindex>
+// CHECK: %[[LIN_IDX:.+]] = arith.addi %[[SPLAT]], %[[IDX]] : vector<8xindex>
+// CHECK: %[[COLLAPSE_I:.+]] = arith.index_cast {{.*}} : index to i64
+// CHECK: %[[VEC:.+]] = xegpu.load %[[COLLAPSE_I]]{{\[}}%[[LIN_IDX]]{{\]}}, %[[MASK]] : i64, vector<8xindex>, vector<8xi1> -> vector<8xf32>
+// CHECK: %[[RES:.+]] = arith.select %[[MASK]], %[[VEC]], %[[PASS_THRU]] : vector<8xi1>, vector<8xf32>
+// CHECK: gpu.return %[[RES]] : vector<8xf32>
+}
diff --git a/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
index b9dd241990222..f20208b299a2d 100644
--- a/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
@@ -204,3 +204,26 @@ gpu.func @scatter_into_subview(%vals: vector<8xf16>,
// CHECK: xegpu.store %[[VALS]], %[[BASE_I64]]{{\[}}%[[LIN]]{{\]}}, %[[MASK]] : vector<8xf16>, i64, vector<8xindex>, vector<8xi1>
// CHECK: gpu.return
}
+
+// -----
+gpu.module @xevm_module {
+gpu.func @store_1D_vector_i32_indices(%vec: vector<8xf32>,
+ %source: memref<8x16x32xf32>, %off1: index, %off2: index, %off3: index,
+ %indices: vector<8xi32>, %mask: vector<8xi1>) {
+ vector.scatter %source[%off1, %off2, %off3][%indices], %mask, %vec
+ : memref<8x16x32xf32>, vector<8xi32>, vector<8xi1>, vector<8xf32>
+ gpu.return
+}
+// Indices that are not index typed are cast before the offset arithmetic.
+// CHECK-LABEL: @store_1D_vector_i32_indices(
+// CHECK-SAME: %[[VALS:.+]]: vector<8xf32>,
+// CHECK-SAME: %[[SRC:.+]]: memref<8x16x32xf32>,
+// CHECK-SAME: %[[OFF1:.+]]: index, %[[OFF2:.+]]: index, %[[OFF3:.+]]: index,
+// CHECK-SAME: %[[INDICES:.+]]: vector<8xi32>, %[[MASK:.+]]: vector<8xi1>) {
+// CHECK: %[[IDX:.+]] = arith.index_cast %[[INDICES]] : vector<8xi32> to vector<8xindex>
+// CHECK: %[[SPLAT:.+]] = vector.broadcast {{.*}} : index to vector<8xindex>
+// CHECK: %[[LIN:.+]] = arith.addi %[[SPLAT]], %[[IDX]] : vector<8xindex>
+// CHECK: %[[BASE_I64:.+]] = arith.index_cast {{.*}} : index to i64
+// CHECK: xegpu.store %[[VALS]], %[[BASE_I64]]{{\[}}%[[LIN]]{{\]}}, %[[MASK]] : vector<8xf32>, i64, vector<8xindex>, vector<8xi1>
+// CHECK: gpu.return
+}
More information about the Mlir-commits
mailing list