[Mlir-commits] [mlir] [mlir][XeGPU] Fix crash in wg-to-sg type converter on non-XeGPU tensors (PR #183914)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/183914

The SCF structural type conversion in XeGPUWgToSgDistributePass registered a RankedTensorType conversion callback that unconditionally called VectorType::get() on the tensor's shape. If the tensor had dynamic dimensions (e.g. tensor<?xi32>) and no XeGPU layout encoding, getSgShapeAndCount() returned the original shape intact (including the kDynamic sentinel value), causing VectorType::get() to abort because VectorType does not support dynamic sizes.

Fix by checking whether the RankedTensorType carries an XeGPU LayoutAttr encoding before attempting the conversion. Plain tensors without such an encoding are left unchanged (std::nullopt causes the passthrough converter to handle them).

Add a regression test for the no-encoding / dynamic-tensor case.

Fixes #182999

>From b178950f17469b945bf57edcb9fa542349a9ccb9 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Feb 2026 05:06:51 -0800
Subject: [PATCH] [mlir][XeGPU] Fix crash in wg-to-sg type converter on
 non-XeGPU tensors

The SCF structural type conversion in XeGPUWgToSgDistributePass
registered a RankedTensorType conversion callback that unconditionally
called VectorType::get() on the tensor's shape. If the tensor had
dynamic dimensions (e.g. tensor<?xi32>) and no XeGPU layout encoding,
getSgShapeAndCount() returned the original shape intact (including the
kDynamic sentinel value), causing VectorType::get() to abort because
VectorType does not support dynamic sizes.

Fix by checking whether the RankedTensorType carries an XeGPU LayoutAttr
encoding before attempting the conversion. Plain tensors without such an
encoding are left unchanged (std::nullopt causes the passthrough
converter to handle them).

Add a regression test for the no-encoding / dynamic-tensor case.

Fixes #182999
---
 .../Transforms/XeGPUWgToSgDistribute.cpp      | 12 +++++++---
 mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir   | 23 +++++++++++++++++++
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 54b5281c37446..3c65091c3b10f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1683,14 +1683,20 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
     converter.addConversion(
         [&](RankedTensorType type,
             SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {
+          // Only convert RankedTensorTypes that carry an XeGPU layout encoding.
+          // Plain tensors (e.g. tensor<?xi32>) have no XeGPU encoding and must
+          // not be converted: VectorType does not support dynamic dimensions.
+          auto encoding =
+              dyn_cast_if_present<xegpu::LayoutAttr>(type.getEncoding());
+          if (!encoding)
+            return std::nullopt;
+
           Type elemTy = type.getElementType();
           ArrayRef<int64_t> shape = type.getShape();
 
           int count;
           SmallVector<int64_t> subShape;
-          std::tie(subShape, count) = getSgShapeAndCount(
-              shape,
-              dyn_cast_if_present<xegpu::LayoutAttr>(type.getEncoding()));
+          std::tie(subShape, count) = getSgShapeAndCount(shape, encoding);
 
           auto newTy = VectorType::get(subShape, elemTy);
           result.append(count, newTy);
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 467c53fa20f94..767113a043557 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -374,3 +374,26 @@ gpu.module @test_1_1_assignment {
     gpu.return
   }
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/182999:
+// Tensors without an XeGPU layout encoding (e.g. tensor<?xi32>) used to crash
+// VectorType::get() because dynamic sizes are not valid for VectorType.
+// The pass should now leave such tensors unchanged.
+
+// CHECK-LABEL: func.func @no_crash_on_dynamic_tensor
+// CHECK-SAME: %[[ARG0:.*]]: tensor<?xi32>
+// CHECK: scf.for
+// CHECK: tensor.insert
+// CHECK: return
+func.func @no_crash_on_dynamic_tensor(%arg0: tensor<?xi32>, %arg1: index) -> tensor<?xi32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %result = scf.for %i = %c0 to %arg1 step %c1 iter_args(%acc = %arg0) -> (tensor<?xi32>) {
+    %val = arith.index_cast %i : index to i32
+    %updated = tensor.insert %val into %acc[%i] : tensor<?xi32>
+    scf.yield %updated : tensor<?xi32>
+  }
+  return %result : tensor<?xi32>
+}