buffer-deallocation: consider aliases introduced by arith.select.

[jreiffers]

Summary

Currently, buffer deallocation considers arith.select to be
non-aliasing, which results in deallocs being inserted incorrectly. Since
arith.select doesn't implement any useful interfaces, this change just handles
it explicitly. Eventually this should probably be fixed properly, if this pass
is going to be used long term.

[jreiffers]

Remove accidental --dump-input.

[springerm]

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

124

Maybe put a TODO that we need an interface here to avoid listing ops. (And adding dialect dependencies to Analysis.)

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

5881

This is a bit unfortunate. But I don't see a good alternative without adding a new interface.

CMake file also needs updating.

[jreiffers]

Fix CMakeLists, add a TODO.

[jreiffers]

Thanks!

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

124

Done.

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

5881

Done.

[herhut]

I find it very strange that arith.select supports memref in the first place. Are there other operations in the arith dialect that do this?

[springerm]

In D132460#3742831, @herhut wrote:

I find it very strange that arith.select supports memref in the first place. Are there other operations in the arith dialect that do this?

I think this is because arith.select used to be std.select and there was no better place to put it. (Maybe the scf dialect would be an alternative.)

[rriddle]

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

124

What is the timeline for fixing this? We should not be adding dialect dependencies to Analysis. Is this going to be fixed soon? Otherwise, please revert.

[jreiffers]

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

124

Ah, thank you for pinging this. I'd be happy to fix this soon, but there should probably be a little bit of discussion on how to do it first. I don't see any existing interface that quite fits here. It might be possible to generalize BranchOpInterface, but to me it doesn't seem like it's going to work well. We could add a new interface to ControlFlowInterfaces (strawman proposal: ValueBranchOpInterface), but that doesn't seem like a perfect match either. Or we could create a new interface elsewhere (but as far as I can tell, it would be a singleton for now, whatever it would end up being. Maybe something specifically around memory aliasing). Input or better ideas welcome.

[springerm]

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

124

I'd say let's move the entire BufferViewFlowAnalysis to the BufferizationTransforms. We only use it for bufferization, so it would make sense to have it there. And we can have a dependency on the Arith dialect there.

Step 2, we can add a new interface that is somewhat similar to BufferizableOpInterface::getAliasingOpResult, but operating on buffers. E.g., we could call it BufferAliasingOpInterface.

Step 3, no longer use RegionBranchOpInterface etc. in BufferViewFlowAnalysis (which is used incorrectly at the moment), but only use BufferAliasingOpInterface. To that end, we have to implement this interface on a few other ops: scf.for, scf.while, scf.execute_region, memref.subview, maybe a few more. The list of ops that must implement this interface is actually quite small. Also, the default case in the analysis (if an op does not implement the interface) should be "assume that operand and result" are aliasing. Unknown ops are currently treated as "no aliasing", which is incorrect.

[jreiffers]

mlir/lib/Analysis/BufferViewFlowAnalysis.cpp

124

OK, I like that solution. https://reviews.llvm.org/D132928 should do step 1 and resolve the immediate issue here. Thanks!

mlir/lib/Analysis/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/Analysis/BufferViewFlowAnalysis.h"
 
+#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 
 /// Constructs a new alias analysis using the op provided.
 BufferViewFlowAnalysis::BufferViewFlowAnalysis(Operation *op) { build(op); }

 
 /// This function constructs a mapping from values to its immediate
 /// dependencies. It iterates over all blocks, gets their predecessors,
 /// determines the values that will be passed to the corresponding block
 /// arguments and inserts them into the underlying map. Furthermore, it wires
 /// successor regions and branch-like return operations from nested regions.
 void BufferViewFlowAnalysis::build(Operation *op) {
   // Registers all dependencies of the given values.
-  auto registerDependencies = [&](auto values, auto dependencies) {
-    for (auto entry : llvm::zip(values, dependencies))
-      this->dependencies[std::get<0>(entry)].insert(std::get<1>(entry));
+  auto registerDependencies = [&](ValueRange values, ValueRange dependencies) {
+    for (auto [value, dep] : llvm::zip(values, dependencies))
+      this->dependencies[value].insert(dep);
   };
 
   // Add additional dependencies created by view changes to the alias list.
   op->walk([&](ViewLikeOpInterface viewInterface) {
     dependencies[viewInterface.getViewSource()].insert(
         viewInterface->getResult(0));
   });
 
for (Region &region : regionInterface->getRegions()) {
           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()});
+  });
 }

mlir/lib/Analysis/CMakeLists.txt

add_mlir_library(MLIRAnalysis
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
 
   DEPENDS
   mlir-headers
 
   LINK_LIBS PUBLIC
+  MLIRArithmeticDialect
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRDataLayoutInterfaces
   MLIRInferIntRangeInterface
   MLIRInferTypeOpInterface
   MLIRLoopLikeInterface
   MLIRSideEffectInterfaces
   MLIRViewLikeInterface
   )
 
 add_subdirectory(Presburger)

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

^bb0(%arg1: memref<?xf32>, %arg2: memref<?xf32>, %arg3: memref<?xf32>):
     scf.yield %q, %b, %arg2: memref<?xf32>, memref<?xf32>, memref<?xf32>
   }
 // CHECK-DAG: memref.dealloc %[[WHILE]]#0
 // CHECK-DAG: memref.dealloc %[[WHILE]]#1
 // CHECK-DAG: memref.dealloc %[[WHILE]]#2
 // CHECK-NEXT: return
   return
 }
+
+// -----
+
+func.func @select_aliases(%arg0: index, %arg1: memref<?xi8>, %arg2: i1) {
+  // CHECK: memref.alloc
+  // CHECK: memref.alloc
+  // CHECK: arith.select
+  // CHECK: test.copy
+  // CHECK: memref.dealloc
+  // CHECK: memref.dealloc
+  %0 = memref.alloc(%arg0) : memref<?xi8>
+  %1 = memref.alloc(%arg0) : memref<?xi8>
+  %2 = arith.select %arg2, %0, %1 : memref<?xi8>
+  test.copy(%2, %arg1) : (memref<?xi8>, memref<?xi8>)
+  return
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     hdrs = glob(
         [
             "include/mlir/Analysis/*.h",
             "include/mlir/Analysis/*/*.h",
         ],
     ),
     includes = ["include"],
     deps = [
+        ":ArithmeticDialect",
         ":CallOpInterfaces",
         ":ControlFlowInterfaces",
         ":DataLayoutInterfaces",
         ":IR",
         ":InferIntRangeInterface",
         ":LoopLikeInterface",
         ":SideEffectInterfaces",
         ":Support",