[Mlir-commits] [mlir] [mlir][xegpu] Cast non-index gather/scatter indices before offset math (PR #216492)

Sat Aug 15 07:13:31 PDT 2026

llvmorg-github-actions[bot] wrote:




@llvm/pr-subscribers-mlir

Author: Hamza Qureshi (hamzaqureshi5)

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/216492.diff


3 Files Affected:

- (modified) mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (+8) 
- (modified) mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir (+26) 
- (modified) mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir (+23) 


``````````diff

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
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/216492
