[Mlir-commits] [mlir] [mlir][shape] Fix crash when folding tensor.extract(shape_of(memref)) (PR #186270)

[Jacques Pienaar]

================
@@ -45,8 +48,11 @@ def TensorCastConstShape : Pat <
   (Tensor_CastOp:$res (Shape_ConstShapeOp $arg)), (Shape_ConstShapeOp $arg),
   [(HasStaticShape $res)]>;
 
-// tensor.extract from shape_of -> tensor.dim. We can take the first index
-// because shape_of always returns a 1D tensor.
+// tensor.extract from shape_of(tensor) -> tensor.dim. We can take the first
+// index because shape_of always returns a 1D tensor.  Only applies when $arg
+// is a tensor; the memref case is not handled here (memref.dim would be
+// needed, but adding a MemRef dependency to this file is not desirable).
 def ExtractFromShapeOfExtentTensor : Pat<
   (Tensor_ExtractOp (Shape_ShapeOfOp $arg), $indices),
----------------
jpienaar wrote:

I thought one could do a Tensor:$arg here to do the same, but could be mixing things up.

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