[mlir] Add canonicalization from `tensor_cast` to `dim` op.

[mravishankar]

Summary

Fold a tensor_cast -> dim to take the dim of the original tensor.

[rriddle]

nit: Can you remove Linalg from the tags?

[mravishankar]

In D93492#2461577, @rriddle wrote:

nit: Can you remove Linalg from the tags?

I dont see where the tag is.

[rriddle]

In D93492#2461595, @mravishankar wrote:

In D93492#2461577, @rriddle wrote:

nit: Can you remove Linalg from the tags?

I dont see where the tag is.

You have [mlir][Linalg] in the title of the commit.

[mravishankar]

In D93492#2461598, @rriddle wrote:

In D93492#2461595, @mravishankar wrote:

In D93492#2461577, @rriddle wrote:

nit: Can you remove Linalg from the tags?

I dont see where the tag is.

You have [mlir][Linalg] in the title of the commit.

:D Done.

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

LogicalResult matchAndRewrite(DimOp dim,
     // Place the load directly after the reshape to ensure that the shape memref
     // was not mutated.
     rewriter.setInsertionPointAfter(reshape);
     rewriter.replaceOpWithNewOp<LoadOp>(dim, reshape.shape(),
                                         llvm::makeArrayRef({dim.index()}));
     return success();
   }
 };
+
+/// Fold dim of a dim of a cast into the the dim of the source of the tensor
+/// cast.
+template <typename CastOpTy>
+struct DimOfCastOp : public OpRewritePattern<DimOp> {
+  using OpRewritePattern<DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp dimOp,
+                                PatternRewriter &rewriter) const override {
+    auto castOp = dimOp.memrefOrTensor().getDefiningOp<CastOpTy>();
+    if (!castOp)
+      return failure();
+    Value newSource = castOp.getOperand();
+    rewriter.replaceOpWithNewOp<DimOp>(dimOp, newSource, dimOp.index());
+    return success();
+  }
+};
+
 } // end anonymous namespace.
 
 void DimOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                         MLIRContext *context) {
-  results.insert<DimOfMemRefReshape>(context);
+  results.insert<DimOfMemRefReshape, DimOfCastOp<TensorCastOp>>(context);
 }
 
 // ---------------------------------------------------------------------------
 // DmaStartOp
 // ---------------------------------------------------------------------------
 
 void DmaStartOp::build(OpBuilder &builder, OperationState &result,
                        Value srcMemRef, ValueRange srcIndices, Value destMemRef,

mlir/test/Dialect/Standard/canonicalize.mlir

func @dim_of_memref_reshape(%arg0: memref<*xf32>, %arg1: memref<?xindex>)
   %c3 = constant 3 : index
   %0 = memref_reshape %arg0(%arg1)
       : (memref<*xf32>, memref<?xindex>) -> memref<*xf32>
   // Update the shape to test that he load ends up in the right place.
   store %c3, %arg1[%c3] : memref<?xindex>
   %1 = dim %0, %c3 : memref<*xf32>
   return %1 : index
 }
+
+// Test case: Folding dim(tensor_cast %0, %idx) -> dim %0, %idx
+// CHECK-LABEL: func @fold_dim_of_tensor_cast
+//  CHECK-SAME:   %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x?xf32>
+//   CHECK-DAG:   %[[C1:.+]] = constant 1 : index
+//   CHECK-DAG:   %[[C4:.+]] = constant 4 : index
+//       CHECK:   %[[T0:.+]] = dim %[[ARG0]], %[[C1]]
+//  CHECK-NEXT:   return %[[C4]], %[[T0]]
+func @fold_dim_of_tensor_cast(%arg0 : tensor<4x?xf32>) -> (index, index) {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %0 = tensor_cast %arg0 : tensor<4x?xf32> to tensor<?x?xf32>
+  %1 = dim %0, %c0 : tensor<?x?xf32>
+  %2 = dim %0, %c1 : tensor<?x?xf32>
+  return %1, %2: index, index
+}