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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

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


2 Files Affected:

- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp (+9-3) 
- (modified) mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir (+23) 


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

``````````

</details>


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