[Mlir-commits] [mlir] [mlir][shape] Fix crash in ShapeOfOpToConstShapeOp (PR #180737)

[Longsheng Mou]

https://github.com/CoTinker updated https://github.com/llvm/llvm-project/pull/180737

>From e51ecea8308be1158b535556b8b093e87247b45c Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 10 Feb 2026 21:46:24 +0800
Subject: [PATCH 1/2] [mlir][shape] Fix crash in ShapeOfOpToConstShapeOp

This PR fixes a crash when `shape.shape_of` op has
static arg and shape result type.
---
 mlir/lib/Dialect/Shape/IR/Shape.cpp       | 10 +++++++---
 mlir/test/Dialect/Shape/canonicalize.mlir | 11 +++++++++++
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index c1210eef4e589..5db18ae5b3b99 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1706,14 +1706,18 @@ struct ShapeOfOpToConstShapeOp : public OpRewritePattern<shape::ShapeOfOp> {
     auto type = llvm::dyn_cast<ShapedType>(op.getArg().getType());
     if (!type || !type.hasStaticShape())
       return failure();
+    Type resultType = op.getResult().getType();
+    if (isa<ShapeType>(resultType))
+      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);
+    if (constShape.getType() != resultType)
+      constShape =
+          tensor::CastOp::create(rewriter, loc, resultType, constShape);
     rewriter.replaceOp(op, constShape);
     return success();
   }
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index f3c25b8c8100e..697bfb19e70a4 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1626,3 +1626,14 @@ func.func @shape_of_0d(%arg0: tensor<f32>) -> tensor<?xindex> {
   %0 = shape.shape_of %arg0 : tensor<f32> -> tensor<?xindex>
   return %0 : tensor<?xindex>
 }
+
+// -----
+
+// Ensure this case not crash.
+
+// CHECK-LABEL: func @shape_of_static_with_shape_result(
+func.func @shape_of_static_with_shape_result(%arg0: tensor<f32>) -> !shape.shape {
+  // CHECK: shape.shape_of
+  %0 = shape.shape_of %arg0 : tensor<f32> -> !shape.shape
+  return %0 : !shape.shape
+}

>From 4494e9cbaca4fa55d58e595da982fad30e4eb235 Mon Sep 17 00:00:00 2001
From: Longsheng Mou <longshengmou at gmail.com>
Date: Tue, 10 Feb 2026 22:53:28 +0800
Subject: [PATCH 2/2] address shape type

---
 mlir/lib/Dialect/Shape/IR/Shape.cpp       | 10 ++++++----
 mlir/test/Dialect/Shape/canonicalize.mlir |  9 ++++-----
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/mlir/lib/Dialect/Shape/IR/Shape.cpp b/mlir/lib/Dialect/Shape/IR/Shape.cpp
index 5db18ae5b3b99..4d03b7b2b2064 100644
--- a/mlir/lib/Dialect/Shape/IR/Shape.cpp
+++ b/mlir/lib/Dialect/Shape/IR/Shape.cpp
@@ -1706,13 +1706,15 @@ struct ShapeOfOpToConstShapeOp : public OpRewritePattern<shape::ShapeOfOp> {
     auto type = llvm::dyn_cast<ShapedType>(op.getArg().getType());
     if (!type || !type.hasStaticShape())
       return failure();
-    Type resultType = op.getResult().getType();
-    if (isa<ShapeType>(resultType))
-      return failure();
 
+    Type resultType = op.getResult().getType();
     Location loc = op.getLoc();
+    Type constResType =
+        isa<ShapeType>(resultType)
+            ? resultType
+            : RankedTensorType::get({type.getRank()}, rewriter.getIndexType());
     Value constShape =
-        ConstShapeOp::create(rewriter, loc,
+        ConstShapeOp::create(rewriter, loc, constResType,
                              rewriter.getIndexTensorAttr(type.getShape()))
             .getResult();
     if (constShape.getType() != resultType)
diff --git a/mlir/test/Dialect/Shape/canonicalize.mlir b/mlir/test/Dialect/Shape/canonicalize.mlir
index 697bfb19e70a4..22add87ff3ed4 100644
--- a/mlir/test/Dialect/Shape/canonicalize.mlir
+++ b/mlir/test/Dialect/Shape/canonicalize.mlir
@@ -1629,11 +1629,10 @@ func.func @shape_of_0d(%arg0: tensor<f32>) -> tensor<?xindex> {
 
 // -----
 
-// Ensure this case not crash.
-
 // CHECK-LABEL: func @shape_of_static_with_shape_result(
-func.func @shape_of_static_with_shape_result(%arg0: tensor<f32>) -> !shape.shape {
-  // CHECK: shape.shape_of
-  %0 = shape.shape_of %arg0 : tensor<f32> -> !shape.shape
+func.func @shape_of_static_with_shape_result(%arg0: tensor<3xf32>) -> !shape.shape {
+  // CHECK: %[[const:.*]] = shape.const_shape [3] : !shape.shape
+  // CHECK: return %[[const]] : !shape.shape
+  %0 = shape.shape_of %arg0 : tensor<3xf32> -> !shape.shape
   return %0 : !shape.shape
 }