[PATCH] D78484: Providing buffer assignment for MLIR

Mon Apr 20 20:36:32 PDT 2020

mehdi_amini added a comment.

Thanks for upstreaming this!

================
Comment at: mlir/include/mlir/Transforms/BufferAssignment.h:9
+//
+// This header file defines buffer assginment helper methods to compute proper
+// and valid positions for placing Alloc and Dealloc operations.
----------------
Typo `assginment`

================
Comment at: mlir/include/mlir/Transforms/BufferAssignment.h:50
+  /// Computes the actual position to place allocs for the given value.
+  OpBuilder::InsertPoint computeAllocPosition(Value value);
+
----------------
There seems to be a precondition that is does not work for a BlockArgument, can you take an `OpResult` as argument instead?

================
Comment at: mlir/include/mlir/Transforms/BufferAssignment.h:58
+/// Helper conversion pattern that encapsulates a BufferAssignmentPlacer
+/// instance.
+template <typename SourceOp>
----------------
Can you expand on the intended uses of it? Maybe a small example?

================
Comment at: mlir/include/mlir/Transforms/BufferAssignment.h:107
+    auto numReturnValues = returnOp.getNumOperands();
+    auto funcOp = returnOp.template getParentOfType<FuncOp>();
+    auto numFuncArgs = funcOp.getNumArguments();
----------------
rriddle wrote:
> Do we need to hard-code FuncOp here? Could we instead just check the number of arguments in the entry block of the function? Or is there some better way to make this more general?
I agree. Some tricky part is that this goes with the other pattern above which is also hard-coded on FuncOp and this pattern can only operate correctly after the other one rewrote the function. This is an even stronger point of coupling I believe.

================
Comment at: mlir/include/mlir/Transforms/BufferAssignment.h:113
+    for (auto operand : llvm::enumerate(operands)) {
+      auto returnArgNumber = numFuncArgs - numReturnValues + operand.index();
+      auto dstBuffer = funcOp.getArgument(returnArgNumber);
----------------
Nit: `numFuncArgs - numReturnValues` is a loop invariant, and it could help readability to hoist it out and name it, like `int firstReturnParameter = numFuncArgs - numReturnValues;` or similar.

================
Comment at: mlir/include/mlir/Transforms/BufferAssignment.h:120
+      rewriter.setInsertionPoint(
+          returnOp.getOperation()->getBlock()->getTerminator());
+      rewriter.create<CopyOpTy>(loc, operand.value(),
----------------
I am puzzled about why not just `rewriter.setInsertionPoint(returnOp.getOperation());` here?

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:44
+// is currently missing is a high-level loop analysis that allows us to move
+// allocs and deallocs outside of the loop blocks.
+//
----------------
When you say "does not support loops": can you be more precise in what will happen in the presence of loops (miscompile, crash, or just suboptimal allocations?)

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:117
+            auto successorOps =
+                branchInterface.getSuccessorOperands(successorIndex);
+            if (successorOps.hasValue()) {
----------------
`successorOps`->`successorOperands` ("Op" is for "Operation" in general)

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:120
+              // Build the actual mapping of values to their immediate aliases.
+              for (auto arg : block.getArguments()) {
+                Value predecessorArgValue =
----------------
Nit: `Value arg`

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:122
+                Value predecessorArgValue =
+                    successorOps.getValue()[arg.getArgNumber()];
+                aliases[predecessorArgValue].insert(arg);
----------------
I'd consider iterating on  `llvm::zip(block.getArguments(), successorOps.getValue())`

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:126
+            }
+          }
+        }
----------------
what happens in the "else" case with respect to aliasing? It isn't clear to me that this is safe?

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:185
+                                       value.getDefiningOp());
+    }
+    // Get all possible aliases
----------------
Nit: avoid trivial braces

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:194
+  /// information.
+  DeallocSetT findAssociatedDeallocs(AllocOp alloc) const {
+    DeallocSetT result;
----------------
Isn't there a trait we could use instead of hard-coding the `AllocOp` here (and everywhere)?

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:232
+                            const DominatorT &doms) const {
+    assert(!value.isa<BlockArgument>() && "Cannot place a block argument");
+    // Start with the current block the value is defined in.
----------------
Can you write the API as taking an `OpResult` instead of the base `Value`?

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:249
+  template <typename AliasesT>
+  Operation *getAllocPosition(Value value, const AliasesT &aliases) const {
+    // Determine the actual block to place the alloc and get liveness
----------------
Does it need to be a template? This is a private method and I see a single use?

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:368
+      }
+    });
+  };
----------------
Isn't this all assuming a particular model with a single allocation / deallocation per buffer?
What about a pattern like the following with two allocs feeding into one dealloc? (we can imagine more complex cases)

```
    cond_br %cond, ^bb1, ^bb2
  ^bb1:
    %buffer1 = alloc ...
    br ^exit(%buffer1 : ....)
  ^bb2:
    %buffer2 = alloc ...
    br ^exit(%buffer2 : ...)
  ^exit(%buffer: ...):
    dealloc %buffer
```

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:388
+  return opBuilder.saveInsertionPoint();
+}
+
----------------
This does not seem to be "computing" anything? Is this like "future work"?

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:409
+    conversion.addInputs(argType.index(), toMemrefConverter(argType.value()));
+  for (auto resType : funcType.getResults())
+    conversion.addInputs(toMemrefConverter(resType));
----------------
Nit: auto -> Type

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:416
+  });
+  // Converting tensor-type block arugments of all blocks inside the
+  // function region to memref-type except for the entry block.
----------------
Typo: `arugments`

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:426
+      oldArg.replaceAllUsesWith(newArg);
+      block.eraseArgument(i + 1);
+    }
----------------
Can't you just call ` setType()` on the BlockArgument?
(also don't use auto when unnecessary)

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:446
+    bool legality = true;
+    for (auto &block2 : funcOp.getBlocks()) {
+      legality &= llvm::all_of(block2.getArguments(), isLegalBlockArg);
----------------
`block` instead of `block2` here?
Also don't use auto when it does not make the code mode readable, I think that `Block &block` would be just fine here.

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:449
+      if (!legality)
+        break;
+    }
----------------
Remove the `legality` boolean and just `if (all_of(...)) return false;` here.

You can even replace the outer loop with an all_of/any_of ;)

```
return all_of(funcOp.getBlocks(), [] (Block &block) {
  return all_of(block2.getArguments(), isLegalBlockArg);
}
```

================
Comment at: mlir/lib/Transforms/BufferAssignment.cpp:464
+      "Executes buffer assignment pass to automatically move alloc and dealloc "
+      "operations into their proper positions");
+}
----------------
It seems to me that the pass is a misnomer: it does not really "assign" buffers, but optimizes the placement, would there be a more accurate name?

Also saying "into their proper positions" seems like it is intended for correctness.

================
Comment at: mlir/test/Transforms/buffer-assignment.mlir:1
+// RUN: mlir-opt -buffer-assignment -allow-unregistered-dialect -split-input-file %s | FileCheck %s -dump-input-on-failure
+
----------------
Can you try to use only operations with registered dialects? The test dialect in-tree is a potential host for such operations, and `std.call @external_func` is also frequently a good escape hatch.

================
Comment at: mlir/test/Transforms/buffer-assignment.mlir:64
+// CHECK-NEXT: dealloc %[[ALLOC]]
+// CHECK-NEXT: return
+
----------------
Can you document the property you're trying to test in each test-case? I really helps to interpret minimal check lines.

================
Comment at: mlir/test/lib/Transforms/TestBufferAssignment.cpp:29
+/// TestBufferAssignmentPreparationPass is equal to
+/// Linalg-Tensor-To-Linalg-Buffer pass and can be separately published.
+/// TODO(dfki): It can be removed later on.
----------------
Can you clarify this sentence? I am missing context

================
Comment at: mlir/tools/mlir-opt/mlir-opt.cpp:98
 void registerTestPasses() {
+  registerBufferAssignmentPass();
   registerConvertToTargetEnvPass();
----------------
This isn't in the test directory, so it shouldn't be registered here but use the same mechanism as the other non-tests passes.

Repository:
  rG LLVM Github Monorepo

CHANGES SINCE LAST ACTION
  https://reviews.llvm.org/D78484/new/

https://reviews.llvm.org/D78484