[mlir] Removed tight coupling of BufferPlacement pass to Alloc and Dealloc operations by using MemoryEffectOpInterface queries.

[dfki-mako]

Summary

The current BufferPlacement implementation tries to find Alloc and Dealloc operations in order to move them. However, this is a tight coupling to standard-dialect ops which has been removed in this CL.

[dfki-mako]

Removed todo comment.

[herhut]

Please wait for @mehdi_amini to take a look.

mlir/lib/Transforms/BufferPlacement.cpp

360

Would just return work?

379–380

Nit: Error messages start lower case.

[pifon2a]

mlir/lib/Transforms/BufferPlacement.cpp

195

nit: add period after "interface".

360

not sure if return would work, but instead of

  if (allocateEffectInstance == effects.end())
        return WalkResult::advance();
  auto allocResult = allocateEffectInstance->getValue().cast<OpResult>();
  placements.emplace_back(
         allocResult, analysis.computeAllocAndDeallocPositions(allocResult));
  return WalkResult::advance();
});

try doing

if (allocateEffectInstance != effects.end()) {
   auto allocResult = allocateEffectInstance->getValue().cast<OpResult>();
   placements.emplace_back(allocResult, 
     analysis.computeAllocAndDeallocPositions(allocResult));
}});

to get rid of WalkResult::advance() .

395–396

I am a bit confused by these two lines:

OpBuilder builder(alloc.getOwner());
builder.setInsertionPointAfter(positions.getDeallocPosition());

the first line creates a builder and sets its insertion point, the second line changes insertion point again. Why not do that in "one hop"? I think i am just confused.

[rriddle]

mlir/lib/Transforms/BufferPlacement.cpp

194–206

This should still be Operation* and not OpOperand?

200

Please make the comments complete sentences.

360

The walk methods are allowed to have void return, so you can strip out the WalkResult::advance. You only need that if you want to interrupt.

[mehdi_amini]

Thanks!
LGTM with the other comments addressed

mlir/lib/Transforms/BufferPlacement.cpp

365

Technically I we could have ops allocating two results, seems like you would only see the first?
We may miscompile in such cases I think, can we not try to move such ops if we see them?

[dfki-mako]

Addressed reviewer comments.

[dfki-mako]

mlir/lib/Transforms/BufferPlacement.cpp

365

Makes sense to me.

[herhut]

Thanks, LGTM!

mlir/lib/Transforms/BufferPlacement.cpp

computeAllocAndDeallocPositions(OpResult result) const {
     // Get all possible aliases.
     auto possibleValues = aliases.resolve(result);
     return BufferPlacementPositions(getAllocPosition(result, possibleValues),
                                     getDeallocPosition(result, possibleValues));
   }
 
   /// Finds all associated dealloc nodes for the alloc nodes using alias
   /// information.
-  DeallocSetT findAssociatedDeallocs(AllocOp alloc) const {
+  DeallocSetT findAssociatedDeallocs(OpResult allocResult) const {
     DeallocSetT result;
-    auto possibleValues = aliases.resolve(alloc);
+    auto possibleValues = aliases.resolve(allocResult);
     for (Value alias : possibleValues)
-      for (Operation *user : alias.getUsers()) {
-        if (isa<DeallocOp>(user))
-          result.insert(user);
+      for (Operation *op : alias.getUsers()) {
+        // Check for an existing memory effect interface.
+        auto effectInstance = dyn_cast<MemoryEffectOpInterface>(op);
+        if (!effectInstance)
+          continue;
+        // Check whether the associated value will be freed using the current
+        // operation.
+        SmallVector<MemoryEffects::EffectInstance, 2> effects;
+        effectInstance.getEffectsOnValue(alias, effects);
+        if (llvm::any_of(effects, [=](MemoryEffects::EffectInstance &it) {
+              return isa<MemoryEffects::Free>(it.getEffect());
+            }))
+          result.insert(op);
       }
     return result;
   }
 
   /// Dumps the buffer placement information to the given stream.
   void print(raw_ostream &os) const {
     os << "// ---- Buffer Placement -----\n";
 

 };
 
 //===----------------------------------------------------------------------===//
 // BufferPlacementPass
 //===----------------------------------------------------------------------===//
 
 /// The actual buffer placement pass that moves alloc and dealloc nodes into
 /// the right positions. It uses the algorithm described at the top of the file.
-// TODO: create a templated version that allows to match dialect-specific
-// alloc/dealloc nodes and to insert dialect-specific dealloc node.
 struct BufferPlacementPass
     : mlir::PassWrapper<BufferPlacementPass, FunctionPass> {
   void runOnFunction() override {
     // Get required analysis information first.
     auto &analysis = getAnalysis<BufferPlacementAnalysis>();
 
     // Compute an initial placement of all nodes.
-    llvm::SmallDenseMap<Value, BufferPlacementPositions, 16> placements;
-    getFunction().walk([&](AllocOp alloc) {
-      placements[alloc] = analysis.computeAllocAndDeallocPositions(
-          alloc.getOperation()->getResult(0));
-      return WalkResult::advance();
+    llvm::SmallVector<std::pair<OpResult, BufferPlacementPositions>, 16>
+        placements;
+    getFunction().walk([&](MemoryEffectOpInterface op) {
+      // Try to find a single allocation result.
+      SmallVector<MemoryEffects::EffectInstance, 2> effects;
+      op.getEffects(effects);
+
+      SmallVector<MemoryEffects::EffectInstance, 2> allocateResultEffects;
+      llvm::copy_if(effects, std::back_inserter(allocateResultEffects),
+                    [=](MemoryEffects::EffectInstance &it) {
+                      Value value = it.getValue();
+                      return isa<MemoryEffects::Allocate>(it.getEffect()) &&
+                             value && value.isa<OpResult>();
+                    });
+      // If there is one result only, we will be able to move the allocation and
+      // (possibly existing) deallocation ops.
+      if (allocateResultEffects.size() == 1) {
+        // Insert allocation result.
+        auto allocResult = allocateResultEffects[0].getValue().cast<OpResult>();
+        placements.emplace_back(
+            allocResult, analysis.computeAllocAndDeallocPositions(allocResult));
+      }
     });
 
-    // Move alloc (and dealloc - if any) nodes into the right places
-    // and insert dealloc nodes if necessary.
-    getFunction().walk([&](AllocOp alloc) {
+    // Move alloc (and dealloc - if any) nodes into the right places and insert
+    // dealloc nodes if necessary.
+    for (auto &entry : placements) {
       // Find already associated dealloc nodes.
+      OpResult alloc = entry.first;
       auto deallocs = analysis.findAssociatedDeallocs(alloc);
       if (deallocs.size() > 1) {
         emitError(alloc.getLoc(),
-                  "Not supported number of associated dealloc operations");
-        return WalkResult::interrupt();
+                  "not supported number of associated dealloc operations");
+        return;
       }
 
       // Move alloc node to the right place.
-      BufferPlacementPositions &positions = placements[alloc];
-      Operation *allocOperation = alloc.getOperation();
+      BufferPlacementPositions &positions = entry.second;
+      Operation *allocOperation = alloc.getOwner();
       allocOperation->moveBefore(positions.getAllocPosition());
 
       // If there is an existing dealloc, move it to the right place.
-      if (deallocs.size()) {
-        Operation *nextOp = positions.getDeallocPosition()->getNextNode();
-        assert(nextOp && "Invalid Dealloc operation position");
+      if (deallocs.size()) {
         (*deallocs.begin())->moveBefore(nextOp);
       } else {
         // If there is no dealloc node, insert one in the right place.
-        OpBuilder builder(alloc);
-        builder.setInsertionPointAfter(positions.getDeallocPosition());
+        OpBuilder builder(nextOp);
         builder.create<DeallocOp>(allocOperation->getLoc(), alloc);
       }
-      return WalkResult::advance();
-    });
+    }
   };
 };
 
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // BufferAssignmentPlacer
 //===----------------------------------------------------------------------===//