[MemRef] Restrict dim(reshape) transform to constant indices

[bkramer]

Summary

Since the load is inserted in place of the reshape (so the shape tensor
cannot be overwritten) the index has to dominate the reshape. With a
dynamic index there I don't see a way of enforcing that. Constants are
moved to the top so it's safe to do this transformation for them.

[pifon2a]

mlir/test/Dialect/Standard/canonicalize.mlir

147

maybe insert a couple of CHECK-NOTs to make sure that it actually folded? memref.dim alone would pass with the original IR.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

struct DimOfMemRefReshape : public OpRewritePattern<DimOp> {
 
   LogicalResult matchAndRewrite(DimOp dim,
                                 PatternRewriter &rewriter) const override {
     auto reshape = dim.memrefOrTensor().getDefiningOp<ReshapeOp>();
 
     if (!reshape)
       return failure();
 
+    // We need to make sure that the index dominates the new load. Since we
+    // can't check that here, restrict the transformation to constant indices.
+    if (!matchPattern(dim.index(), matchConstantIndex()))
+      return failure();
+
     // Place the load directly after the reshape to ensure that the shape memref
     // was not mutated.
     rewriter.setInsertionPointAfter(reshape);
     Location loc = dim.getLoc();
     Value load = rewriter.create<LoadOp>(loc, reshape.shape(), dim.index());
     if (load.getType() != dim.getType())
       load = rewriter.create<IndexCastOp>(loc, dim.getType(), load);
     rewriter.replaceOp(dim, load);

mlir/test/Dialect/Standard/canonicalize.mlir

func @dim_of_memref_reshape_i32(%arg0: memref<*xf32>, %arg1: memref<?xi32>)
   %0 = memref.reshape %arg0(%arg1)
       : (memref<*xf32>, memref<?xi32>) -> memref<*xf32>
   %1 = memref.dim %0, %c3 : memref<*xf32>
   return %1 : index
 }
 
 // -----
 
+// Test case: Folding of memref.dim(memref.reshape %v %shp, %idx) -> memref.load %shp[%idx]
+// CHECK-LABEL: func @dim_of_memref_reshape_not_dominating(
+//   CHECK:   memref.dim
+func @dim_of_memref_reshape_not_dominating(%arg0: memref<*xf32>, %arg1: memref<?xi32>, %arg2: memref<?xindex>)
+    -> index {
+  %c3 = constant 3 : index
+  %0 = memref.reshape %arg0(%arg1)
+      : (memref<*xf32>, memref<?xi32>) -> memref<*xf32>
+  %idx = memref.load %arg2[%c3] : memref<?xindex>
+  %1 = memref.dim %0, %idx : memref<*xf32>
+  return %1 : index
+}
+
+// -----
+
 // Test case: Folding memref.dim(tensor.cast %0, %idx) -> memref.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:.+]] = memref.dim %[[ARG0]], %[[C1]]
 //  CHECK-NEXT:   return %[[C4]], %[[T0]]
 func @fold_dim_of_tensor.cast(%arg0 : tensor<4x?xf32>) -> (index, index) {