[mlir] Canonicalize AllocOp's with only store and dealloc uses

[Hardcode84]

clang-format

[mehdi_amini]

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

198–201

206

This loop could be llvm::any_of I think

[Hardcode84]

review comments

[bondhugula]

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

198–201

Canonicalizatoin of alloc's that only had dealloc's was already supported. This is why the comment here mentioned "no uses" - because the SimplifyDeadDealloc would get rid of all deallocs if those had been the only things besides alloc's, and then the alloc's would have zero uses and disappear. You can remove the SimplifyDeadDealloc pattern now?

[Hardcode84]

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

198–201

SimplifyDeadDealloc already removed

[Hardcode84]

@bondhugula ping

[Hardcode84]

ping

[bondhugula]

In D100268#2694177, @Hardcode84 wrote:

@bondhugula ping

I'm sorry this revision never showed up in my "Waiting for review" list for some reason. I did make multiple passes from time to time on my "waiting for my review", but didn't find this.

[bondhugula]

LGTM. Please note the remaining minor comments.

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

198–201

with only store and dealloc uses -> with no uses or only store and dealloc uses

?

209

auto -> Operation

[Hardcode84]

rebase, rewiew comments

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     ```mlir
     %0 = memref.alloc() : memref<8x64xf32, (d0, d1) -> (d0, d1), 1>
     memref.dealloc %0 : memref<8x64xf32, (d0, d1) -> (d0, d1), 1>
     ```
   }];
 
   let arguments = (ins Arg<AnyMemRef, "", [MemFree]>:$memref);
 
-  let hasCanonicalizer = 1;
   let hasFolder = 1;
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//
 

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

LogicalResult matchAndRewrite(AllocLikeOp alloc,
     auto resultCast =
         rewriter.create<CastOp>(alloc.getLoc(), newAlloc, alloc.getType());
 
     rewriter.replaceOp(alloc, {resultCast});
     return success();
   }
 };
 
-/// Fold alloc operations with no uses. Alloc has side effects on the heap,
-/// but can still be deleted if it has zero uses.
-struct SimplifyDeadAlloc : public OpRewritePattern<AllocOp> {
-  using OpRewritePattern<AllocOp>::OpRewritePattern;
+/// Fold alloc operations with no users or only store and dealloc uses.
+template <typename T>
+struct SimplifyDeadAlloc : public OpRewritePattern<T> {
+  using OpRewritePattern<T>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(AllocOp alloc,
+  LogicalResult matchAndRewrite(T alloc,
                                 PatternRewriter &rewriter) const override {
-    if (alloc.use_empty()) {
-      rewriter.eraseOp(alloc);
-      return success();
-    }
-    return failure();
-  }
+    if (llvm::any_of(alloc->getUsers(), [](Operation *op) {
+          return !isa<StoreOp, DeallocOp>(op);
+        }))
+      return failure();
+
+    for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
+      rewriter.eraseOp(user);
+
 };
 } // end anonymous namespace.
 
 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                           MLIRContext *context) {
-  results.add<SimplifyAllocConst<AllocOp>, SimplifyDeadAlloc>(context);
+  results.add<SimplifyAllocConst<AllocOp>, SimplifyDeadAlloc<AllocOp>>(context);
 }
 
 void AllocaOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                            MLIRContext *context) {
-  results.add<SimplifyAllocConst<AllocaOp>>(context);
+  results.add<SimplifyAllocConst<AllocaOp>, SimplifyDeadAlloc<AllocaOp>>(
+      context);
 }
 
 //===----------------------------------------------------------------------===//
 // AssumeAlignmentOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(AssumeAlignmentOp op) {
   unsigned alignment = op.alignment();

 
 OpFoldResult CloneOp::fold(ArrayRef<Attribute> operands) {
   return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
-namespace {
-/// Fold Dealloc operations that are deallocating an AllocOp that is only used
-/// by other Dealloc operations.
-struct SimplifyDeadDealloc : public OpRewritePattern<DeallocOp> {
-  using OpRewritePattern<DeallocOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(DeallocOp dealloc,
-                                PatternRewriter &rewriter) const override {
-    // Check that the memref operand's defining operation is an AllocOp.
-    Value memref = dealloc.memref();
-    if (!isa_and_nonnull<AllocOp>(memref.getDefiningOp()))
-      return failure();
-
-    // Check that all of the uses of the AllocOp are other DeallocOps.
-    for (auto *user : memref.getUsers())
-      if (!isa<DeallocOp>(user))
-        return failure();
-
-    // Erase the dealloc operation.
-    rewriter.eraseOp(dealloc);
-    return success();
-  }
-};
-} // end anonymous namespace.
 
 static LogicalResult verify(DeallocOp op) {
   if (!op.memref().getType().isa<MemRefType>())
     return op.emitOpError("operand must be a memref");
   return success();
 }
 
-void DeallocOp::getCanonicalizationPatterns(RewritePatternSet &results,
-                                            MLIRContext *context) {
-  results.add<SimplifyDeadDealloc>(context);
-}
-
 LogicalResult DeallocOp::fold(ArrayRef<Attribute> cstOperands,
                               SmallVectorImpl<OpFoldResult> &results) {
   /// dealloc(memrefcast) -> dealloc
   return foldMemRefCast(*this);
 }
 
 //===----------------------------------------------------------------------===//
 // DimOp

mlir/test/Dialect/Affine/canonicalize.mlir

func @compose_affine_maps_1dto2d_with_symbols() {
   return
 }
 
 // -----
 
 // CHECK-DAG: #[[$MAP8:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
 // CHECK-DAG: #[[$MAP8a:.*]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
 
-// CHECK-LABEL: func @compose_affine_maps_2d_tile() {
-func @compose_affine_maps_2d_tile() {
-  %0 = memref.alloc() : memref<16x32xf32>
-  %1 = memref.alloc() : memref<16x32xf32>
-
+// CHECK-LABEL: func @compose_affine_maps_2d_tile
+func @compose_affine_maps_2d_tile(%0: memref<16x32xf32>, %1: memref<16x32xf32>) {
   %c4 = constant 4 : index
   %c8 = constant 8 : index
 
   affine.for %i0 = 0 to 3 {
     %x0 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i0)[%c4]
     affine.for %i1 = 0 to 3 {
       %x1 = affine.apply affine_map<(d0)[s0] -> (d0 ceildiv s0)> (%i1)[%c8]
       affine.for %i2 = 0 to 3 {
func @compose_affine_maps_multiple_symbols(%arg0: index, %arg1: index) -> index {
   %e = affine.apply affine_map<(d0)[s0] -> (d0 floordiv s0)> (%c)[%arg1]
   // CHECK: [[I0:.*]] = affine.apply #[[$MAP14]]()[%{{.*}}, %{{.*}}]
   return %e : index
 }
 
 // -----
 
 // CHECK-LABEL: func @arg_used_as_dim_and_symbol
-func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) {
+func @arg_used_as_dim_and_symbol(%arg0: memref<100x100xf32>, %arg1: index, %arg2: f32) -> (memref<100x100xf32, 1>, memref<1xi32>) {
   %c9 = constant 9 : index
   %1 = memref.alloc() : memref<100x100xf32, 1>
   %2 = memref.alloc() : memref<1xi32>
   affine.for %i0 = 0 to 100 {
     affine.for %i1 = 0 to 100 {
       %3 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 + s0 + s1)>
         (%i0, %i1)[%arg1, %c9]
       %4 = affine.apply affine_map<(d0, d1, d3) -> (d3 - (d0 + d1))>
         (%arg1, %c9, %3)
       // CHECK: memref.store %arg2, %{{.*}}[%{{.*}}, %{{.*}}]
       memref.store %arg2, %1[%4, %arg1] : memref<100x100xf32, 1>
     }
   }
-  return
+  return %1, %2 : memref<100x100xf32, 1>, memref<1xi32>
 }
 
 // -----
 
 // CHECK-LABEL: func @trivial_maps
 func @trivial_maps() {
   // CHECK-NOT: affine.apply
 

mlir/test/Dialect/Tensor/canonicalize.mlir

func @extract_from_tensor.generate_2d(%idx0: index, %idx1: index, %tensor: tensor<*xf32>) -> index {
   // CHECK-NEXT: return %[[RES]]
   return %4 : index
 }
 
 // -----
 
 // CHECK-LABEL: func @extract_from_tensor.generate_sideeffects
 // CHECK-SAME: %[[IDX:.*]]: index
-func @extract_from_tensor.generate_sideeffects(%idx: index, %tensor: tensor<*xf32>) -> index {
+func @extract_from_tensor.generate_sideeffects(%idx: index, %tensor: tensor<*xf32>, %mem: memref<?xindex>) -> index {
   %size = rank %tensor : tensor<*xf32>
-  %mem = memref.alloc(%size) : memref<?xindex>
   // CHECK: %[[DTENSOR:.*]] = tensor.generate
   %0 = tensor.generate %size {
     ^bb0(%arg0: index):
     %1 = memref.dim %tensor, %arg0 : tensor<*xf32>
     memref.store %1, %mem[%arg0] : memref<?xindex>
     tensor.yield %1 : index
   } : tensor<?xindex>
   // CHECK: %[[RES:.*]] = tensor.extract %[[DTENSOR]][%[[IDX]]]

mlir/test/Transforms/canonicalize.mlir

^bb1:
   memref.dealloc %a: memref<4xf32>
   return
 
 ^bb2:
   memref.dealloc %a: memref<4xf32>
   return
 }
 
+// CHECK-LABEL: func @write_only_alloc_fold
+func @write_only_alloc_fold(%v: f32) {
+  // CHECK-NEXT: return
+  %c0 = constant 0 : index
+  %c4 = constant 4 : index
+  %a = memref.alloc(%c4) : memref<?xf32>
+  memref.store %v, %a[%c0] : memref<?xf32>
+  memref.dealloc %a: memref<?xf32>
+  return
+}
+
+// CHECK-LABEL: func @write_only_alloca_fold
+func @write_only_alloca_fold(%v: f32) {
+  // CHECK-NEXT: return
+  %c0 = constant 0 : index
+  %c4 = constant 4 : index
+  %a = memref.alloca(%c4) : memref<?xf32>
+  memref.store %v, %a[%c0] : memref<?xf32>
+  return
+}
+
 // CHECK-LABEL: func @dead_block_elim
 func @dead_block_elim() {
   // CHECK-NOT: ^bb
   func @nested() {
     return
 
   ^bb1:
     return
   }
   return
 
 ^bb1:
   return
 }
 
 // CHECK-LABEL: func @dyn_shape_fold(%arg0: index, %arg1: index)
-func @dyn_shape_fold(%L : index, %M : index) -> (memref<? x ? x i32>, memref<? x ? x f32>, memref<4 x ? x 8 x ? x ? x f32>) {
+func @dyn_shape_fold(%L : index, %M : index) -> (memref<4 x ? x 8 x ? x ? x f32>, memref<? x ? x i32>, memref<? x ? x f32>, memref<4 x ? x 8 x ? x ? x f32>) {
   // CHECK: %c0 = constant 0 : index
   %zero = constant 0 : index
   // The constants below disappear after they propagate into shapes.
   %nine = constant 9 : index
   %N = constant 1024 : index
   %K = constant 512 : index
 
-  // CHECK-NEXT: memref.alloc(%arg0) : memref<?x1024xf32>
+  // CHECK: memref.alloc(%arg0) : memref<?x1024xf32>
   %a = memref.alloc(%L, %N) : memref<? x ? x f32>
 
-  // CHECK-NEXT: memref.alloc(%arg1) : memref<4x1024x8x512x?xf32>
+  // CHECK: memref.alloc(%arg1) : memref<4x1024x8x512x?xf32>
   %b = memref.alloc(%N, %K, %M) : memref<4 x ? x 8 x ? x ? x f32>
 
-  // CHECK-NEXT: memref.alloc() : memref<512x1024xi32>
+  // CHECK: memref.alloc() : memref<512x1024xi32>
   %c = memref.alloc(%K, %N) : memref<? x ? x i32>
 
   // CHECK: memref.alloc() : memref<9x9xf32>
   %d = memref.alloc(%nine, %nine) : memref<? x ? x f32>
 
   // CHECK: memref.alloca(%arg1) : memref<4x1024x8x512x?xf32>
   %e = memref.alloca(%N, %K, %M) : memref<4 x ? x 8 x ? x ? x f32>
 
   // CHECK: affine.for
   affine.for %i = 0 to %L {
     // CHECK-NEXT: affine.for
     affine.for %j = 0 to 10 {
       // CHECK-NEXT: memref.load %0[%arg2, %arg3] : memref<?x1024xf32>
       // CHECK-NEXT: memref.store %{{.*}}, %1[%c0, %c0, %arg2, %arg3, %c0] : memref<4x1024x8x512x?xf32>
       %v = memref.load %a[%i, %j] : memref<?x?xf32>
       memref.store %v, %b[%zero, %zero, %i, %j, %zero] : memref<4x?x8x?x?xf32>
     }
   }
 
-  return %c, %d, %e : memref<? x ? x i32>, memref<? x ? x f32>, memref<4 x ? x 8 x ? x ? x f32>
+  return %b, %c, %d, %e : memref<4 x ? x 8 x ? x ? x f32>, memref<? x ? x i32>, memref<? x ? x f32>, memref<4 x ? x 8 x ? x ? x f32>
 }
 
 #map1 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 #map2 = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s2 + d1 * s1 + d2 + s0)>
 #map3 = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 
 // CHECK-LABEL: func @dim_op_fold(
 // CHECK-SAME: %[[ARG0:[a-z0-9]*]]: index