[MLIR][XLA] Buffer Assignment

[dfki-ehna]

In this PR, we have provided our first draft for generic XLA buffer assignment. This generic BufferAssignment class automatically analyzes the values and their aliases (also in other blocks) and returns the positions of allocations and deallocations. To find these positions, the algorithm uses the block Dominator and Post-Dominator analyses. In our proposed algorithm, we have considered aliasing, liveness analysis, nested regions, branches, conditional branches, critical edges, and independency to custom block terminators. This implementation doesn't support block loops. However, we have considered this in our design. For this purpose, we only need to have a loop analysis to allocate and deallocate some special cases outside of these loops.

This is the sample usage of BufferAssignment class:

BufferAssignment assignments(funcOp);
auto positions = assignments.computeAllocAndDeallocPositions(value);
allocBuilder.setInsertionPoint(positions.getAllocPosition());
<create alloc>
deallocBuilder.setInsertionPointAfter(positions.getDeallocPosition());
<create dealloc>

Or alternatively:

BufferAssignment assignments(funcOp);
auto positions = assignments.computeAllocAndDeallocPositions(value);
positions.insertAlloc<AllocOp>(...);
positions.insertDealloc<DeallocOp>(...);

Please note that this class should be used during the legalization (i.e. HLO-To-LHLO).

[pifon2a]

First portion of nits. More coming.

[dfki-ehna]

Refactored BufferAssignment analysis into two distinct parts:

• BufferAssignmentLegalizer helper class that can be used during legalization.
• BufferAssignmentPass that moves alloc and dealloc nodes (if any) into the
  right positions. If there are not associated dealloc nodes for a given
  alloc node, this pass inserts a DeallocOp in the right place automatically.

[sherhut]

Could you provide an example how the use of this would look like for hlo->lhlo? It is not clear to me how this composes.

[dfki-ehna]

The intended use of this buffer assignment is that the dialect-expert creates an instance of BufferAssignmentLegalizer at the beginning of the pass (i.e. in the runOnFunction()) and passes this instance to all conversion-pattern classes using OwningRewritePatternList instance.
These conversion pattern classes could inherit from BufferAssignmentOpConversionPattern that has a protected BufferAssignmentLegalizer member. This will avoid creating BufferAssignmentLegalizer multiple times. In each conversion pattern class that the dialect-expert needs to allocate a buffer, they do the following:

BufferAssignmentOpConversionPattern<HloOpTy>::bufferAssignment->computeAllocPosition(result);
auto alloc = allocPosition.insertAlloc<AllocOp>(result, loc, memref_type); // we are going to remove result from the list of arguments in the next commit since it's unnecessary.

Later on, --buffer-assignment flag must be passed to tf-opt to move Alloc and Dealloc operations in their proper positions. If there is no Dealloc for a buffer, it will be automatically inserted.

[sherhut]

Let's get this landed and improve further in tree. Thanks!

[sherhut]

Thanks!

[joker-eph]

Is there a test covering this function? (our test coverage shows it uncovered)

[joker-eph]

I'm missing where is this pass covered by tests?

[joker-eph]

Lacking my comment being addressed, I issued a revert here: #38291

[byronyi]

This unfortunately leaves this PR to a terminal state. @dfki-ehna would you mind to open a new one? Thanks!

[dfki-mako]

@joker-eph Oh, I see 🔢 Maybe we can add a test dialect including a legalization step in order to simulate lowering from one dialect to another one (or even the same) with buffers?

[joker-eph]

Why do we need a test dialect? I may be missing something but the description of the pass is:

 /// The actual buffer assignment pass that moves alloc and dealloc nodes into
 /// the right positions. 

I was expecting this to be able to run on something like LHLO?

[pifon2a]

@joker-eph The test dialect will be needed if we want to have this pass in MLIR core, since there is no LHLO there. I think @dfki-mako was assuming this.

[dfki-mako]

@joker-eph I have already thought about this transfer and future features at the same time. It might be useful to have a small test dialect that contains additional features that we might want to think about and include in the general BA pass.

[dfki-mako]

@joker-eph Please note that this is not a strict requirement at the moment. We can include several Linalg-based tests in the MLIR core version and apply BA to them.

[joker-eph]

I am a bit lost (but it is quite late here, my brain may be slow...): I was looking at this PR in TensorFlow, in a disconnected way from Core, and I see code that does not seem to have any test here.
Regardless of plans about upstreaming, I expect that the development in MLIR-TensorFlow follows roughly the same set of good practices as upstream, starting with complete coverage of new code with lit tests. Unless I'm mistaken there is a non-tested pass in this PR?

[dfki-mako]

Yes and no. We had to redesign several parts and split some features (regarding legalization and use) into other PRs. No, because the test pass basically computes all required information to move allocs and deallocs. Yes, because the actual movement part is not covered.

[dfki-ehna]

@joker-eph you are right with respect to this PR. It has never been our intention to push an untested pass. We address the missing-tests issue and add a complete test case to test the whole BA pass including movement of allocs and deallocs.

[joker-eph]

We address the missing-tests issue and add a complete test case to test the whole BA pass including movement of allocs and deallocs.

Is there already a PR up for review?