[Mlir-commits] [mlir] [MLIR][Shape] Fix crash in ShapeOfOp canonicalization during inlining (PR #178839)

[llvmlistbot at llvm.org]

https://github.com/Sejal-07661 updated https://github.com/llvm/llvm-project/pull/178839

>From ca3a51ac7178ded3e182c8456433598bbdfd41c2 Mon Sep 17 00:00:00 2001
From: Sejal-07661 <sejalsomkuwar4 at gmail.com>
Date: Fri, 30 Jan 2026 10:40:33 +0530
Subject: [PATCH 1/2] [MLIR][Shape] Fix crash in ShapeOfOp canonicalization
 during inlining

---
 mlir/lib/Dialect/Shape/IR/Shape.cpp       | 33 ++++++++++++-----------
 mlir/test/Dialect/Shape/canonicalize.mlir | 18 +++++++++++++
 2 files changed, 36 insertions(+), 15 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 95546b4a4deb7..a3aa87adc5b30 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1700,21 +1700,24 @@ struct ShapeOfOpToConstShapeOp : public OpRewritePattern<shape::ShapeOfOp> {
   using OpRewritePattern<shape::ShapeOfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(shape::ShapeOfOp op,
-                                PatternRewriter &rewriter) const override {
-    auto type = llvm::dyn_cast<ShapedType>(op.getArg().getType());
-    if (!type || !type.hasStaticShape())
-      return failure();
-    Location loc = op.getLoc();
-    Value constShape =
-        ConstShapeOp::create(rewriter, loc,
-                             rewriter.getIndexTensorAttr(type.getShape()))
-            .getResult();
-    if (constShape.getType() != op.getResult().getType())
-      constShape = tensor::CastOp::create(rewriter, loc,
-                                          op.getResult().getType(), constShape);
-    rewriter.replaceOp(op, constShape);
-    return success();
-  }
+                              PatternRewriter &rewriter) const override {
+  auto shapedTy = dyn_cast<ShapedType>(op.getArg().getType());
+  if (!shapedTy || !shapedTy.hasStaticShape())
+    return failure();
+
+  Location loc = op.getLoc();
+  auto attr = rewriter.getIndexTensorAttr(shapedTy.getShape());
+
+  // IMPORTANT: Build const_shape with the SAME result type as shape_of.
+  // This avoids creating a tensor.cast to a non-tensor type (!shape.shape).
+  Value constShape = rewriter
+      .create<shape::ConstShapeOp>(loc, op.getType(), attr)
+      .getResult();
+
+  rewriter.replaceOp(op, constShape);
+  return success();
+}
+
 };
 
 // Canonicalize
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 32a7c0186f282..1e1ef4b131e3c 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1614,3 +1614,21 @@ func.func @shape_of_0d(%arg0: tensor<f32>) -> tensor<?xindex> {
   %0 = shape.shape_of %arg0 : tensor<f32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
+
+
+// --------------------------------------------
+// ---- shape.shape_of canonicalize regression ----
+
+// RUN: mlir-opt %s -canonicalize | FileCheck %s
+
+module {
+  func.func @main(%arg0: tensor<2x3xi32>) -> !shape.shape {
+    %0 = shape.shape_of %arg0 : tensor<2x3xi32> -> !shape.shape
+    return %0 : !shape.shape
+  }
+}
+
+// CHECK-LABEL: func.func @main
+// CHECK:       shape.const_shape [2, 3] : !shape.shape
+// CHECK-NOT:   shape.shape_of
+

>From f87fd6da35fcb40917e5f9fb66b2df4f47097541 Mon Sep 17 00:00:00 2001
From: Sejal-07661 <sejalsomkuwar4 at gmail.com>
Date: Tue, 10 Feb 2026 01:32:31 +0530
Subject: [PATCH 2/2] [MLIR][Shape] Fix ShapeOfOp canonicalization for dynamic
 extent tensors

---
 mlir/lib/Dialect/Shape/IR/Shape.cpp       | 51 ++++++++++++++++-------
 mlir/test/Dialect/Shape/canonicalize.mlir | 25 +++++------
 2 files changed, 48 insertions(+), 28 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index a3aa87adc5b30..739c7bf5ea427 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1700,26 +1700,49 @@ struct ShapeOfOpToConstShapeOp : public OpRewritePattern<shape::ShapeOfOp> {
   using OpRewritePattern<shape::ShapeOfOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(shape::ShapeOfOp op,
-                              PatternRewriter &rewriter) const override {
-  auto shapedTy = dyn_cast<ShapedType>(op.getArg().getType());
-  if (!shapedTy || !shapedTy.hasStaticShape())
-    return failure();
+                                PatternRewriter &rewriter) const override {
+    auto shapedTy = dyn_cast<ShapedType>(op.getArg().getType());
+    if (!shapedTy || !shapedTy.hasStaticShape())
+      return failure();
 
-  Location loc = op.getLoc();
-  auto attr = rewriter.getIndexTensorAttr(shapedTy.getShape());
+    Location loc = op.getLoc();
+    auto shapeAttr = rewriter.getIndexTensorAttr(shapedTy.getShape());
+    Type resultTy = op.getType();
 
-  // IMPORTANT: Build const_shape with the SAME result type as shape_of.
-  // This avoids creating a tensor.cast to a non-tensor type (!shape.shape).
-  Value constShape = rewriter
-      .create<shape::ConstShapeOp>(loc, op.getType(), attr)
-      .getResult();
+    // Case 1: Result is !shape.shape
+    if (llvm::isa<shape::ShapeType>(resultTy)) {
+      Value cs =
+          rewriter.create<shape::ConstShapeOp>(loc, resultTy, shapeAttr)
+              .getResult();
+      rewriter.replaceOp(op, cs);
+      return success();
+    }
 
-  rewriter.replaceOp(op, constShape);
-  return success();
-}
+    // Case 2: Result is an extent tensor (tensor<...xindex>)
+    auto resultTensorTy = llvm::dyn_cast<RankedTensorType>(resultTy);
+    if (!resultTensorTy)
+      return failure();
 
+    // Build const_shape with the concrete ranked type implied by the constant.
+    int64_t rank = static_cast<int64_t>(shapedTy.getShape().size());
+    auto concreteTy = RankedTensorType::get({rank}, rewriter.getIndexType());
+
+    Value cs =
+        rewriter.create<shape::ConstShapeOp>(loc, concreteTy, shapeAttr)
+            .getResult();
+
+    // If shape_of wanted tensor<?xindex>, cast from tensor<Rxindex>.
+    if (concreteTy != resultTensorTy) {
+      cs = rewriter.create<tensor::CastOp>(loc, resultTensorTy, cs).getResult();
+    }
+
+    rewriter.replaceOp(op, cs);
+    return success();
+  }
 };
 
+
+
 // Canonicalize
 //
 // %0 = tensor.reshape %input(%shape) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 1e1ef4b131e3c..4b929d6e0cbeb 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1616,19 +1616,16 @@ func.func @shape_of_0d(%arg0: tensor<f32>) -> tensor<?xindex> {
 }
 
 
-// --------------------------------------------
-// ---- shape.shape_of canonicalize regression ----
-
-// RUN: mlir-opt %s -canonicalize | FileCheck %s
-
-module {
-  func.func @main(%arg0: tensor<2x3xi32>) -> !shape.shape {
-    %0 = shape.shape_of %arg0 : tensor<2x3xi32> -> !shape.shape
-    return %0 : !shape.shape
-  }
-}
+// -----
 
-// CHECK-LABEL: func.func @main
+// CHECK-LABEL: func.func @shape_of_canonicalize_regression
 // CHECK:       shape.const_shape [2, 3] : !shape.shape
-// CHECK-NOT:   shape.shape_of
-
+// CHECK:       shape.const_shape [2, 3] : tensor<2xindex>
+// CHECK:       tensor.cast
+func.func @shape_of_canonicalize_regression(%arg0: tensor<2x3xi32>)
+    -> (!shape.shape, tensor<2xindex>, tensor<?xindex>) {
+  %0 = shape.shape_of %arg0 : tensor<2x3xi32> -> !shape.shape
+  %1 = shape.shape_of %arg0 : tensor<2x3xi32> -> tensor<2xindex>
+  %2 = shape.shape_of %arg0 : tensor<2x3xi32> -> tensor<?xindex>
+  return %0, %1, %2 : !shape.shape, tensor<2xindex>, tensor<?xindex>
+}
\ No newline at end of file