[MLIR][memref] Fix findDealloc() to handle > 1 dealloc for the given alloc.

[jurahul]

Summary

• Change findDealloc() to return Optional<Operation *> and return None if > 1 dealloc is associated with the given alloc.
• Add findDeallocs() to return all deallocs associated with the given alloc.
• Fix current uses of findDealloc() to bail out if > 1 dealloc is found.

[mehdi_amini]

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

544

I'd check the first before calling the second findDealloc: early exit before doing useless work :)

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

39

When I saw the .h, I thought you were keeping this API for efficiency. But seeing the implementation it clearly isn't :)

Is this separate API really worth it?

[jurahul]

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

39

All current uses actually use the second API. May be I should remove the first API (it can be added later when someone needs to look at all deallocations) and make the second one efficient.

[jurahul]

Remove findDeallocs()

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

let description = [{
     of the clone operation after the clone operation thus leads to undefined
     behavior.
   }];
 
   let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "", []>:$input);
   let results = (outs Arg<AnyRankedOrUnrankedMemRef, "", []>:$output);
 
   let extraClassDeclaration = [{
-    Value getSource() { return input();}
+    Value getSource() { return input(); }
     Value getTarget() { return output(); }
   }];
 
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 
   let hasFolder = 1;
   let verifier = ?;
   let hasCanonicalizer = 1;

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

 
 #ifndef MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 #define MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
 namespace mlir {
 
-/// Finds the associated dealloc that can be linked to our allocation nodes (if
-/// any).
-Operation *findDealloc(Value allocValue);
-
+/// Finds a single dealloc operation for the given allocated value. If there
+/// are > 1 deallocates for `allocValue`, returns None, else returns the single
+/// deallocate if it exists or nullptr.
+llvm::Optional<Operation *> findDealloc(Value allocValue);
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H

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

 /// 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(T alloc,
                                 PatternRewriter &rewriter) const override {
     if (llvm::any_of(alloc->getUsers(), [&](Operation *op) {
           if (auto storeOp = dyn_cast<StoreOp>(op))
             return storeOp.value() == alloc;
           return !isa<DeallocOp>(op);
         }))
       return failure();
 
     for (Operation *user : llvm::make_early_inc_range(alloc->getUsers()))
       rewriter.eraseOp(user);
 
     rewriter.eraseOp(alloc);
     return success();
effects.emplace_back(MemoryEffects::Read::get(), input(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), output(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Allocate::get(), output(),
                        SideEffects::DefaultResource::get());
 }
 
 namespace {
-/// Fold Dealloc operations that are deallocating an AllocOp that is only used
-/// by other Dealloc operations.
+/// Merge the clone and its source (by converting the clone to a cast) when
+/// possible.
 struct SimplifyClones : public OpRewritePattern<CloneOp> {
   using OpRewritePattern<CloneOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CloneOp cloneOp,
                                 PatternRewriter &rewriter) const override {
     if (cloneOp.use_empty()) {
       rewriter.eraseOp(cloneOp);
       return success();
     }
 
     Value source = cloneOp.input();
 
     // This only finds dealloc operations for the immediate value. It should
     // also consider aliases. That would also make the safety check below
     // redundant.
-    Operation *cloneDeallocOp = findDealloc(cloneOp.output());
-    Operation *sourceDeallocOp = findDealloc(source);
+    llvm::Optional<Operation *> maybeCloneDeallocOp =
+        findDealloc(cloneOp.output());
+    // Skip if either of them has > 1 deallocate operations.
+    if (!maybeCloneDeallocOp.hasValue())
+      return failure();
+    llvm::Optional<Operation *> maybeSourceDeallocOp = findDealloc(source);
+    if (!maybeSourceDeallocOp.hasValue())
+      return failure();
+    Operation *cloneDeallocOp = *maybeCloneDeallocOp;
+    Operation *sourceDeallocOp = *maybeSourceDeallocOp;
 
     // If both are deallocated in the same block, their in-block lifetimes
     // might not fully overlap, so we cannot decide which one to drop.
     if (cloneDeallocOp && sourceDeallocOp &&
         cloneDeallocOp->getBlock() == sourceDeallocOp->getBlock())
       return failure();
 
     Block *currentBlock = cloneOp->getBlock();

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

-//===- Utils.cpp - Utilities to support the MemRef dialect ----------------===//
+//===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements utilities for the MemRef dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 using namespace mlir;
 
-/// Finds associated deallocs that can be linked to our allocation nodes (if
-/// any).
-Operation *mlir::findDealloc(Value allocValue) {
-  auto userIt = llvm::find_if(allocValue.getUsers(), [&](Operation *user) {
+/// Finds a single dealloc operation for the given allocated value.
+llvm::Optional<Operation *> mlir::findDealloc(Value allocValue) {
+  Operation *dealloc = nullptr;
+  for (Operation *user : allocValue.getUsers()) {
     auto effectInterface = dyn_cast<MemoryEffectOpInterface>(user);
     if (!effectInterface)
-      return false;
+      continue;
     // Try to find a free effect that is applied to one of our values
     // that will be automatically freed by our pass.
     SmallVector<MemoryEffects::EffectInstance, 2> effects;
     effectInterface.getEffectsOnValue(allocValue, effects);
-    return llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) {
+    const bool isFree =
+        llvm::any_of(effects, [&](MemoryEffects::EffectInstance &it) {
           return isa<MemoryEffects::Free>(it.getEffect());
         });
-  });
-  // Assign the associated dealloc operation (if any).
-  return userIt != allocValue.user_end() ? *userIt : nullptr;
+    if (!isFree)
+      continue;
+    // If we found > 1 dealloc, return None.
+    if (dealloc)
+      return llvm::None;
+    dealloc = user;
+  }
+  return dealloc;
 }

mlir/lib/Transforms/BufferUtils.cpp

llvm::copy_if(
         });
     // If there is one result only, we will be able to move the allocation and
     // (possibly existing) deallocation ops.
     if (allocateResultEffects.size() != 1)
       return;
     // Get allocation result.
     Value allocValue = allocateResultEffects[0].getValue();
     // Find the associated dealloc value and register the allocation entry.
-    allocs.push_back(std::make_tuple(allocValue, findDealloc(allocValue)));
+    llvm::Optional<Operation *> dealloc = findDealloc(allocValue);
+    // If the allocation has > 1 dealloc associated with it, skip handling it.
+    if (!dealloc.hasValue())
+      return;
+    allocs.push_back(std::make_tuple(allocValue, *dealloc));
   });
 }
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementTransformationBase
 //===----------------------------------------------------------------------===//
 
 /// Constructs a new transformation base using the given root operation.

mlir/test/Dialect/MemRef/canonicalize.mlir

func @alias_is_freed(%arg0 : memref<?xf32>) {
   memref.dealloc %arg0 : memref<?xf32>
   "use"(%1) : (memref<32xf32>) -> ()
   memref.dealloc %1 : memref<32xf32>
   return
 }
 
 // -----
 
+// Verify SimplifyClones skips clones with multiple deallocations.
+// CHECK-LABEL: @clone_multiple_dealloc_of_source
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+func @clone_multiple_dealloc_of_source(%arg0: memref<?xf32>) -> memref<?xf32> {
+  // CHECK-NEXT: %[[RES:.*]] = memref.clone %[[ARG]]
+  // CHECK: memref.dealloc %[[ARG]]
+  // CHECK: memref.dealloc %[[ARG]]
+  // CHECK: return %[[RES]]
+  %0 = memref.clone %arg0 : memref<?xf32> to memref<?xf32>
+  "if_else"() ({
+    memref.dealloc %arg0 : memref<?xf32>
+    }, {
+    memref.dealloc %arg0 : memref<?xf32>
+    }) : () -> ()
+  return %0 : memref<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @clone_multiple_dealloc_of_clone
+// CHECK-SAME: %[[ARG:.*]]: memref<?xf32>
+func @clone_multiple_dealloc_of_clone(%arg0: memref<?xf32>) -> memref<?xf32> {
+  // CHECK-NEXT: %[[CLONE:.*]] = memref.clone %[[ARG]]
+  // CHECK: memref.dealloc %[[CLONE]]
+  // CHECK: memref.dealloc %[[CLONE]]
+  // CHECK: return %[[ARG]]
+  %0 = memref.clone %arg0 : memref<?xf32> to memref<?xf32>
+  "use"(%0) : (memref<?xf32>) -> ()
+  "if_else"() ({
+    memref.dealloc %0 : memref<?xf32>
+    }, {
+    memref.dealloc %0 : memref<?xf32>
+    }) : () -> ()
+  return %arg0 : memref<?xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @dim_of_sized_view
 //  CHECK-SAME:   %{{[a-z0-9A-Z_]+}}: memref<?xi8>
 //  CHECK-SAME:   %[[SIZE:.[a-z0-9A-Z_]+]]: index
 //       CHECK:   return %[[SIZE]] : index
 func @dim_of_sized_view(%arg : memref<?xi8>, %size: index) -> index {
   %c0 = constant 0 : index
   %0 = memref.reinterpret_cast %arg to offset: [0], sizes: [%size], strides: [0] : memref<?xi8> to memref<?xi8>
   %1 = memref.dim %0, %c0 : memref<?xi8>

 
 // -----
 // CHECK-LABEL: func @allocator
 // CHECK:   %[[alloc:.+]] = memref.alloc
 // CHECK:   memref.store %[[alloc:.+]], %arg0
 func @allocator(%arg0 : memref<memref<?xi32>>, %arg1 : index)  {
   %0 = memref.alloc(%arg1) : memref<?xi32>
   memref.store %0, %arg0[] : memref<memref<?xi32>>
-  return 
+  return
 }
 
 // -----
 
 func @collapsing_memref_reshapes_to_zero_dim(%arg0 : memref<1x1x1xf32>)
                                              -> memref<f32> {
   %0 = memref.collapse_shape %arg0 [[0, 1, 2]]
       : memref<1x1x1xf32> into memref<1xf32>