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

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-03T12:21:23+01:00
New Revision: a8a5242bb2dc2949e148e7000230d479be321def

URL: https://github.com/llvm/llvm-project/commit/a8a5242bb2dc2949e148e7000230d479be321def
DIFF: https://github.com/llvm/llvm-project/commit/a8a5242bb2dc2949e148e7000230d479be321def.diff

LOG: [mlir][XeGPU] Fix crash in wg-to-sg type converter on non-XeGPU tensors (#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

Added: 
    

Modified: 
    mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
    mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir

Removed: 
    


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