[Mlir-commits] [mlir] [mlir][xegpu] Fix crash on 0D vector in vector-to-xegpu transfer lowering (PR #205812)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Chennes (Chennesxu)

<details>
<summary>Changes</summary>

`transferPreconditions` accepted rank-0 vectors, which then reached `computeOffsets` and indexed an empty `SmallVector` (`broadcasted[0]`), asserting `idx < size()` and crashing. For example:

```mlir
vector.transfer_write %arg0, %arg1[%1] : vector<f32>, memref<3xf32>
```

This is valid IR (the verifier accepts it; 0D vector transfers are part of the vector dialect spec), so the conversion pass should not crash on it.

Bail out via `notifyMatchFailure` for 0D vectors in the shared `transferPreconditions`, so both `transfer_read` and `transfer_write` are left unconverted instead of crashing. This matches how the core vector lowering already treats 0D transfers as an unsupported corner case (see `mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp`).

Regression tests are added for 0D `transfer_read` and `transfer_write`.

Fixes #<!-- -->205281


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


3 Files Affected:

- (modified) mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp (+2) 
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir (+18) 
- (modified) mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir (+17) 


``````````diff
diff --git a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
index 9038bf35b6b15..670a4d37b91a1 100644
--- a/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
+++ b/mlir/lib/Conversion/VectorToXeGPU/VectorToXeGPU.cpp
@@ -92,6 +92,8 @@ static LogicalResult transferPreconditions(PatternRewriter &rewriter,
 
   VectorType vecTy = xferOp.getVectorType();
   unsigned vecRank = vecTy.getRank();
+  if (vecRank == 0)
+    return rewriter.notifyMatchFailure(xferOp, "0D vectors are not supported");
   if (xferOp.hasOutOfBoundsDim() && vecRank < 2)
     return rewriter.notifyMatchFailure(
         xferOp, "Boundary check is available only for block instructions.");
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
index daa53c0877c8f..55baffec994a9 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-read-to-xegpu.mlir
@@ -667,3 +667,21 @@ gpu.func @load_1D_vector_alloca_promoted_unsupported(%offset: index)
 // LOAD-GATHER-NOT: xegpu.load_matrix
 
 }
+
+// -----
+gpu.module @xevm_module {
+gpu.func @load_0D_vector_unsupported(%source: memref<3xf32>,
+    %offset: index) -> vector<f32> {
+  %c0 = arith.constant 0.0 : f32
+  %0 = vector.transfer_read %source[%offset], %c0
+    : memref<3xf32>, vector<f32>
+  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
+
+}
diff --git a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
index a7cd3d7652d85..755b8a3abc4e0 100644
--- a/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
+++ b/mlir/test/Conversion/VectorToXeGPU/transfer-write-to-xegpu.mlir
@@ -412,3 +412,20 @@ gpu.func @store_1D_vector_addrspace3_unsupported(%vec: vector<8xf32>,
 // STORE-SCATTER: vector.transfer_write
 
 }
+
+// -----
+gpu.module @xevm_module {
+gpu.func @store_0D_vector_unsupported(%vec: vector<f32>,
+    %source: memref<3xf32>, %offset: index) {
+  vector.transfer_write %vec, %source[%offset]
+    : vector<f32>, memref<3xf32>
+  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
+
+}

``````````

</details>


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