[mlir][bufferization] Fix unknown ops in BufferViewFlowAnalysis

[springerm]

Summary

If an op is unknown to the analysis, it must be treated conservatively: assume that every operand aliases with every result.

mlir/lib/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.cpp

 //======- BufferViewFlowAnalysis.cpp - Buffer alias analysis -*- C++ -*-======//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"
 
-#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/SetOperations.h"
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 
 /// Constructs a new alias analysis using the op provided.

 /// successor regions and branch-like return operations from nested regions.
 void BufferViewFlowAnalysis::build(Operation *op) {
   // Registers all dependencies of the given values.
   auto registerDependencies = [&](ValueRange values, ValueRange dependencies) {
     for (auto [value, dep] : llvm::zip(values, dependencies))
       this->dependencies[value].insert(dep);
   };
 
+  op->walk([&](Operation *op) {
+    // TODO: We should have an op interface instead of a hard-coded list of
+    // interfaces/ops.
+
     // Add additional dependencies created by view changes to the alias list.
-  op->walk([&](ViewLikeOpInterface viewInterface) {
+    if (auto viewInterface = dyn_cast<ViewLikeOpInterface>(op)) {
       dependencies[viewInterface.getViewSource()].insert(
           viewInterface->getResult(0));
-  });
+      return WalkResult::advance();
+    }
 
+    if (auto branchInterface = dyn_cast<BranchOpInterface>(op)) {
       // Query all branch interfaces to link block argument dependencies.
-  op->walk([&](BranchOpInterface branchInterface) {
       Block *parentBlock = branchInterface->getBlock();
       for (auto it = parentBlock->succ_begin(), e = parentBlock->succ_end();
            it != e; ++it) {
         // Query the branch op interface to get the successor operands.
         auto successorOperands =
             branchInterface.getSuccessorOperands(it.getIndex());
         // Build the actual mapping of values to their immediate dependencies.
         registerDependencies(successorOperands.getForwardedOperands(),
                              (*it)->getArguments().drop_front(
                                  successorOperands.getProducedOperandCount()));
       }
-  });
+      return WalkResult::advance();
+    }
 
+    if (auto regionInterface = dyn_cast<RegionBranchOpInterface>(op)) {
       // Query the RegionBranchOpInterface to find potential successor regions.
-  op->walk([&](RegionBranchOpInterface regionInterface) {
       // Extract all entry regions and wire all initial entry successor inputs.
       SmallVector<RegionSuccessor, 2> entrySuccessors;
       regionInterface.getSuccessorRegions(/*index=*/std::nullopt,
                                           entrySuccessors);
       for (RegionSuccessor &entrySuccessor : entrySuccessors) {
         // Wire the entry region's successor arguments with the initial
         // successor inputs.
         assert(entrySuccessor.getSuccessor() &&
                "Invalid entry region without an attached successor region");
         registerDependencies(
             regionInterface.getSuccessorEntryOperands(
                 entrySuccessor.getSuccessor()->getRegionNumber()),
             entrySuccessor.getSuccessorInputs());
       }
 
       // Wire flow between regions and from region exits.
       for (Region &region : regionInterface->getRegions()) {
         // Iterate over all successor region entries that are reachable from the
         // current region.
         SmallVector<RegionSuccessor, 2> successorRegions;
         regionInterface.getSuccessorRegions(region.getRegionNumber(),
                                             successorRegions);
         for (RegionSuccessor &successorRegion : successorRegions) {
           // Determine the current region index (if any).
           std::optional<unsigned> regionIndex;
           Region *regionSuccessor = successorRegion.getSuccessor();
           if (regionSuccessor)
             regionIndex = regionSuccessor->getRegionNumber();
           // Iterate over all immediate terminator operations and wire the
           // successor inputs with the successor operands of each terminator.
           for (Block &block : region) {
             auto successorOperands = getRegionBranchSuccessorOperands(
                 block.getTerminator(), regionIndex);
             if (successorOperands) {
               registerDependencies(*successorOperands,
                                    successorRegion.getSuccessorInputs());
             }
           }
         }
       }
-  });
 
-  // TODO: This should be an interface.
-  op->walk([&](arith::SelectOp selectOp) {
-    registerDependencies({selectOp.getOperand(1)}, {selectOp.getResult()});
-    registerDependencies({selectOp.getOperand(2)}, {selectOp.getResult()});
+      return WalkResult::advance();
+    }
+
+    // Unknown op: Assume that all operands alias with all results.
+    for (Value operand : op->getOperands()) {
+      if (!isa<BaseMemRefType>(operand.getType()))
+        continue;
+      for (Value result : op->getResults()) {
+        if (!isa<BaseMemRefType>(result.getType()))
+          continue;
+        registerDependencies({operand}, {result});
+      }
+    }
+    return WalkResult::advance();
   });
 }

mlir/test/Dialect/Bufferization/Transforms/buffer-deallocation.mlir

func.func @select_aliases(%arg0: index, %arg1: memref<?xi8>, %arg2: i1) {
   %1 = memref.alloc(%arg0) : memref<?xi8>
   %2 = arith.select %arg2, %0, %1 : memref<?xi8>
   test.copy(%2, %arg1) : (memref<?xi8>, memref<?xi8>)
   return
 }
 
 // -----
 
+func.func @f(%arg0: memref<f64>) -> memref<f64> {
+  return %arg0 : memref<f64>
+}
+
+// CHECK-LABEL: func @function_call
+//       CHECK:   memref.alloc
+//       CHECK:   memref.alloc
+//       CHECK:   call
+//       CHECK:   test.copy
+//       CHECK:   memref.dealloc
+//       CHECK:   memref.dealloc
+func.func @function_call() {
+  %alloc = memref.alloc() : memref<f64>
+  %alloc2 = memref.alloc() : memref<f64>
+  %ret = call @f(%alloc) : (memref<f64>) -> memref<f64>
+  test.copy(%ret, %alloc2) : (memref<f64>, memref<f64>)
+  return
+}
+
+// -----
+
 // Memref allocated in `then` region and passed back to the parent if op.
 #set = affine_set<() : (0 >= 0)>
 // CHECK-LABEL:  func @test_affine_if_1
 // CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>) -> memref<10xf32> {
 func.func @test_affine_if_1(%arg0: memref<10xf32>) -> memref<10xf32> {
   %0 = affine.if #set() -> memref<10xf32> {
     %alloc = memref.alloc() : memref<10xf32>
     affine.yield %alloc : memref<10xf32>