[mlir] Canonicalize static alloc followed by memref_cast and std.view

[asaadaldien]

Summary

Rewrite alloc, memref_cast, std.view into allo, std.view by droping memref_cast.

[merge_guards_bot]

Unit tests: pass. 61305 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/lib/Dialect/StandardOps/Ops.cpp

2538

Remove trivial braces.

2546

Can you use named accessors for these instead?

Also, initializer_list is a temporary so this is likely to break. Use an array instead, or inline this into the 'replace' call below.

mlir/test/Transforms/canonicalize.mlir

739

Don't check SSA names directly, they are not preserved.

https://mlir.llvm.org/getting_started/TestingGuide/

[nicolasvasilache]

LGTM conditioned on addressing River's review.

[asaadaldien]

Resolve comments....

[asaadaldien]

mlir/lib/Dialect/StandardOps/Ops.cpp

2546

used ViewOp::operands instead apparently its variadic operands only

[merge_guards_bot]

Unit tests: pass. 61305 tests passed, 0 failed and 736 were skipped.

clang-tidy: fail. Please fix clang-tidy findings.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

[rriddle]

mlir/test/Transforms/canonicalize.mlir

739

You didn't fix this.

[asaadaldien]

mlir/test/Transforms/canonicalize.mlir

739

Oh yes changed different line! will send a diff to use ALLOC_MEM instead of %0 everywhere here!

mlir/lib/Dialect/StandardOps/Ops.cpp

auto newViewOp = rewriter.create<ViewOp>(viewOp.getLoc(), newMemRefType,
                                              viewOp.getOperand(0), newOperands);
     // Insert a cast so we have the same type as the old memref type.
     rewriter.replaceOpWithNewOp<MemRefCastOp>(droppedOperands, viewOp,
                                               newViewOp, viewOp.getType());
     return matchSuccess();
   }
 };
 
+struct ViewOpMemrefCastFolder : public OpRewritePattern<ViewOp> {
+  using OpRewritePattern<ViewOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(ViewOp viewOp,
+                                     PatternRewriter &rewriter) const override {
+    Value memrefOperand = viewOp.getOperand(0);
+    MemRefCastOp memrefCastOp =
+        dyn_cast_or_null<MemRefCastOp>(memrefOperand.getDefiningOp());
+    if (!memrefCastOp)
+      return matchFailure();
+    Value allocOperand = memrefCastOp.getOperand();
+    AllocOp allocOp = dyn_cast_or_null<AllocOp>(allocOperand.getDefiningOp());
+    if (!allocOp)
+      return matchFailure();
+    rewriter.replaceOpWithNewOp<ViewOp>(memrefOperand, viewOp, viewOp.getType(),
+                                        allocOperand, viewOp.operands());
+    return matchSuccess();
+  }
+};
+
 } // end anonymous namespace
 
 void ViewOp::getCanonicalizationPatterns(OwningRewritePatternList &results,
                                          MLIRContext *context) {
-  results.insert<ViewOpShapeFolder>(context);
+  results.insert<ViewOpShapeFolder, ViewOpMemrefCastFolder>(context);
 }
 
 //===----------------------------------------------------------------------===//
 // SubViewOp
 //===----------------------------------------------------------------------===//
 
 // Returns a MemRefType with dynamic sizes and offset and the same stride as the
 // `memRefType` passed as argument.

mlir/test/Transforms/canonicalize.mlir

 // CHECK-DAG: #[[VIEW_MAP0:map[0-9]+]] = (d0, d1) -> (d0 * 11 + d1 + 15)
 // CHECK-DAG: #[[VIEW_MAP1:map[0-9]+]] = (d0, d1)[s0] -> (d0 * 11 + s0 + d1)
 // CHECK-DAG: #[[VIEW_MAP2:map[0-9]+]] = (d0, d1)[s0] -> (d0 * s0 + d1 + 15)
 // CHECK-DAG: #[[VIEW_MAP3:map[0-9]+]] = (d0, d1, d2)[s0] -> (d0 * s0 + d1 * 7 + d2)
 // CHECK-DAG: #[[VIEW_MAP4:map[0-9]+]] = (d0, d1) -> (d0 * 4 + d1 + 15)
 
 // CHECK-LABEL: func @view
 func @view(%arg0 : index) {
+  // CHECK: %[[ALLOC_MEM:.*]] = alloc() : memref<2048xi8>
   %0 = alloc() : memref<2048xi8>
   %c0 = constant 0 : index
   %c7 = constant 7 : index
   %c11 = constant 11 : index
   %c15 = constant 15 : index
 
   // Test: fold constant sizes and offset, update map with static stride/offset.
   // CHECK: std.view %0[][] : memref<2048xi8> to memref<7x11xf32, #[[VIEW_MAP0]]>
func @view(%arg0 : index) {
   // stride on dim 1, but leave dynamic stride on dim 0.
   // CHECK: std.view %0[][%arg0, %arg0] : memref<2048xi8> to memref<?x?x7xf32, #[[VIEW_MAP3]]>
   %4 = view %0[][%arg0, %arg0, %c7]
     : memref<2048xi8> to memref<?x?x?xf32, #TEST_VIEW_MAP1>
   load %4[%c0, %c0, %c0] : memref<?x?x?xf32, #TEST_VIEW_MAP1>
 
   // Test: preserve an existing static dim size while folding a dynamic
   // dimension and offset.
-  // CHECK: std.view %0[][] : memref<2048xi8> to memref<7x4xf32, #[[VIEW_MAP4]]>
-  %5 = view %0[%c15][%c7]
-    : memref<2048xi8> to memref<?x4xf32, #TEST_VIEW_MAP2>
+  // CHECK: std.view %[[ALLOC_MEM]][][] : memref<2048xi8> to memref<7x4xf32, #[[VIEW_MAP4]]>
+  %5 = view %0[%c15][%c7] : memref<2048xi8> to memref<?x4xf32, #TEST_VIEW_MAP2>
   load %5[%c0, %c0] : memref<?x4xf32, #TEST_VIEW_MAP2>
 
+  // Test: folding static alloc and memref_cast into a view.
+  // CHECK: std.view %0[][%c15, %c7] : memref<2048xi8> to memref<?x?xf32>
+  %6 = memref_cast %0 : memref<2048xi8> to memref<?xi8>
+  %7 = view %6[%c15][%c7] : memref<?xi8> to memref<?x?xf32>
+  load %7[%c0, %c0] : memref<?x?xf32>
   return
 }
 
 // -----
 
 // CHECK-DAG: #[[BASE_MAP0:map[0-9]+]] = (d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2)
 // CHECK-DAG: #[[SUBVIEW_MAP0:map[0-9]+]] = (d0, d1, d2)[s0] -> (d0 * 64 + s0 + d1 * 4 + d2)
 // CHECK-DAG: #[[SUBVIEW_MAP1:map[0-9]+]] = (d0, d1, d2) -> (d0 * 64 + d1 * 4 + d2 + 79)