[Mlir-commits] [mlir] [MLIR] Fix canonicalization pattern for 'shape.shape_of' (PR #134234)

[Mehdi Amini]

================
@@ -1734,10 +1734,22 @@ struct ShapeOfFromReshape : public OpRewritePattern<shape::ShapeOfOp> {
     // Operand 'shape' of 'tensor.reshape' may now be used as the result of
     // 'shape.shape_of'. While its type is guaranteed to be compatible in well-
     // formed IR, it may not be identical (dynamically vs statically shaped),
-    // in which case it needs to be cast first.
+    // in which case it needs to be cast first using 'tensor.cast'.
+    // Additionally, it may not have identical element type (i32 vs index) 
+    // while it has identical shaped type (dynamic vs static), in which case it needs 
+    // to be cast first using 'arith.index_cast'.
+    // Note: 'shape.shape_of' op result must be shape or extent tensor.
     Value shape = tensorReshapeOp.getShape();
-    if (op.getType() != shape.getType())
-      shape = rewriter.create<tensor::CastOp>(op.getLoc(), op.getType(), shape);
+
+    auto opTensorType = llvm::dyn_cast<RankedTensorType>(op.getType());
+    auto shapeTensorType = llvm::dyn_cast<RankedTensorType>(shape.getType());
----------------
joker-eph wrote:

If you check for it, then we should have tests showing a crash without the checks though :)

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