[mlir][std] Canonicalize a dim(memref_reshape) into a load from the shape operand

[herhut]

Summary

This canonicalization helps propagate shape information through the program.

[ftynse]

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

1573

Nit: I don't remember if the pattern applicator has an InsertGuard or if the patterns are supposed to preserve the insertion point.

[rriddle]

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

1573

There is no guard, but the applicator always updates the position before applying each pattern so it should be fine.

[frgossen]

Thanks!

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

let builders = [
     OpBuilderDAG<(ins "Value":$memrefOrTensor, "Value":$index)>
   ];
 
   let extraClassDeclaration = [{
     /// Helper function to get the index as a simple integer if it is constant.
     Optional<int64_t> getConstantIndex();
   }];
 
+  let hasCanonicalizer = 1;
   let hasFolder = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DivFOp
 //===----------------------------------------------------------------------===//
 
 def DivFOp : FloatArithmeticOp<"divf"> {

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

OpFoldResult DimOp::fold(ArrayRef<Attribute> operands) {
 
   // dim(memrefcast) -> dim
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
 
   return {};
 }
 
+namespace {
+/// Fold dim of a memref reshape operation to a load into the reshape's shape
+/// operand.
+struct DimOfMemRefReshape : public OpRewritePattern<DimOp> {
+  using OpRewritePattern<DimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DimOp dim,
+                                PatternRewriter &rewriter) const override {
+    auto reshape = dim.memrefOrTensor().getDefiningOp<MemRefReshapeOp>();
+
+    if (!reshape)
+      return failure();
+
+    // 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();
+  }
+};
+} // end anonymous namespace.
+
+void DimOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
+                                        MLIRContext *context) {
+  results.insert<DimOfMemRefReshape>(context);
+}
+
 // ---------------------------------------------------------------------------
 // DmaStartOp
 // ---------------------------------------------------------------------------
 
 void DmaStartOp::build(OpBuilder &builder, OperationState &result,
                        Value srcMemRef, ValueRange srcIndices, Value destMemRef,
                        ValueRange destIndices, Value numElements,
                        Value tagMemRef, ValueRange tagIndices, Value stride,

mlir/test/Dialect/Standard/canonicalize.mlir

func @cmpi_equal_operands(%arg0: i64)
   %5 = cmpi "ne", %arg0, %arg0 : i64
   %6 = cmpi "slt", %arg0, %arg0 : i64
   %7 = cmpi "sgt", %arg0, %arg0 : i64
   %8 = cmpi "ult", %arg0, %arg0 : i64
   %9 = cmpi "ugt", %arg0, %arg0 : i64
   return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9
       : i1, i1, i1, i1, i1, i1, i1, i1, i1, i1
 }
+
+// Test case: Folding of dim(memref_reshape %v %shp, %idx) -> load %shp[%idx]
+// CHECK-LABEL: func @dim_of_memref_reshape(
+//  CHECK-SAME:     %[[MEM:[0-9a-z]+]]: memref<*xf32>,
+//  CHECK-SAME:     %[[SHP:[0-9a-z]+]]: memref<?xindex>
+//  CHECK-NEXT:   %[[IDX:.*]] = constant 3
+//  CHECK-NEXT:   %[[DIM:.*]] = load %[[SHP]][%[[IDX]]]
+//  CHECK-NEXT:   store
+//   CHECK-NOT:   dim
+//       CHECK:   return %[[DIM]] : index
+func @dim_of_memref_reshape(%arg0: memref<*xf32>, %arg1: memref<?xindex>)
+    -> index {
+  %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
+}