[Mlir-commits] [mlir] [mlir][xegpu] Skip scalable gather/scatter lowering (PR #234725)

[Akimasa Watanuki]

https://github.com/Men-cotton created https://github.com/llvm/llvm-project/pull/234725

Leave scalable `vector.gather` and `vector.scatter` operations unconverted because XeGPU load/store offsets and masks require fixed-size vectors. Check this before creating offset arithmetic to avoid mixing scalable and fixed-size vector types and triggering a folding assertion.

Add regression coverage for retaining scalable gathers and scatters.

Fixes #202531

Assisted-by: Codex / GPT-6

>From 79072b6edd0ed8048f94a2e7b7b8e260ff694a46 Mon Sep 17 00:00:00 2001
From: mencotton <mencotton0410 at gmail.com>
Date: Thu, 22 Oct 2026 13:44:59 -0500
Subject: [PATCH] [mlir][xegpu] Skip scalable gather/scatter lowering

Leave scalable vector.gather and vector.scatter operations unconverted because XeGPU load/store offsets and masks require fixed-size vectors. Check this before creating offset arithmetic to avoid mixing scalable and fixed-size vector types and triggering a folding assertion.

Add regression coverage for retaining scalable gathers and scatters.

Fixes #202531

Assisted-by: Codex / GPT-6
---
 .../VectorToXeGPU/VectorToXeGPU.cpp           |  8 +++++++
 .../VectorToXeGPU/gather-to-xegpu.mlir        | 21 +++++++++++++++++++
 .../VectorToXeGPU/scatter-to-xegpu.mlir       | 17 +++++++++++++++
 3 files changed, 46 insertions(+)

diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 6b0c09b3b48f4b..3999d84eaf6a98 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -951,6 +951,10 @@ struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
     if (!srcTy)
       return rewriter.notifyMatchFailure(gatherOp, "Expects memref source");
 
+    if (gatherOp.getVectorType().isScalable())
+      return rewriter.notifyMatchFailure(gatherOp,
+                                         "Scalable vectors are not supported");
+
     Location loc = gatherOp.getLoc();
     VectorType vectorType = gatherOp.getVectorType();
 
@@ -986,6 +990,10 @@ struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
     if (!srcTy)
       return rewriter.notifyMatchFailure(scatterOp, "Expects memref source");
 
+    if (scatterOp.getVectorType().isScalable())
+      return rewriter.notifyMatchFailure(scatterOp,
+                                         "Scalable vectors are not supported");
+
     Location loc = scatterOp.getLoc();
     auto meta = computeMemrefMeta(scatterOp, rewriter);
     if (meta.first.empty())
diff --git a/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
index 9617d347af2b6d..82c41195bea6fc 100644
--- a/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/gather-to-xegpu.mlir
@@ -26,6 +26,27 @@ gpu.func @load_1D_vector(%source: memref<8x16x32xf32>,
 // CHECK:        gpu.return %[[RES]] : vector<8xf32>
 }
 
+// -----
+
+// XeGPU load offsets and masks must be fixed-size vectors.
+gpu.module @xevm_module {
+gpu.func @no_load_scalable_vector(%source: memref<?xf32>,
+    %indices: vector<[3]xindex>, %mask: vector<[3]xi1>,
+    %pass_thru: vector<[3]xf32>) -> vector<[3]xf32> {
+  %c0 = arith.constant 0 : index
+  %0 = vector.gather %source[%c0][%indices], %mask, %pass_thru
+    : memref<?xf32>, vector<[3]xindex>, vector<[3]xi1>, vector<[3]xf32> into vector<[3]xf32>
+  gpu.return %0 : vector<[3]xf32>
+}
+// CHECK-LABEL: @no_load_scalable_vector
+// CHECK-SAME: %[[SRC:.*]]: memref<?xf32>,
+// CHECK-SAME: %[[INDICES:.*]]: vector<[3]xindex>, %[[MASK:.*]]: vector<[3]xi1>,
+// CHECK-SAME: %[[PASS:.*]]: vector<[3]xf32>
+// CHECK: %[[OFF:.*]] = arith.constant 0 : index
+// CHECK: %[[RES:.*]] = vector.gather %[[SRC]][%[[OFF]]] [%[[INDICES]]], %[[MASK]], %[[PASS]]
+// CHECK: gpu.return %[[RES]] : vector<[3]xf32>
+}
+
 // -----
 gpu.module @xevm_module {
 gpu.func @load_2D_memref(%source: memref<8x32xf32>,
diff --git a/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
index b9dd241990222c..6cc53daf70fda3 100644
--- a/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/scatter-to-xegpu.mlir
@@ -22,6 +22,23 @@ gpu.func @store_1D_vector(%vec: vector<8xf32>, %source: memref<8x16x32xf32>,
 // CHECK:        gpu.return
 }
 
+// -----
+
+// XeGPU store offsets and masks must be fixed-size vectors.
+gpu.module @xevm_module {
+gpu.func @no_store_scalable_vector(%value: vector<[3]xf32>,
+    %source: memref<?xf32>, %offset: index,
+    %indices: vector<[3]xindex>, %mask: vector<[3]xi1>) {
+  vector.scatter %source[%offset][%indices], %mask, %value
+    : memref<?xf32>, vector<[3]xindex>, vector<[3]xi1>, vector<[3]xf32>
+  gpu.return
+}
+// CHECK-LABEL: @no_store_scalable_vector
+// CHECK-SAME: %[[VAL:.*]]: vector<[3]xf32>, %[[SRC:.*]]: memref<?xf32>,
+// CHECK-SAME: %[[OFF:.*]]: index, %[[INDICES:.*]]: vector<[3]xindex>, %[[MASK:.*]]: vector<[3]xi1>
+// CHECK: vector.scatter %[[SRC]][%[[OFF]]] [%[[INDICES]]], %[[MASK]], %[[VAL]]
+}
+
 // -----
 gpu.module @xevm_module {
 gpu.func @store_2D_memref(%vec: vector<8xf32>, %source: memref<8x32xf32>,