This is an archive of the discontinued LLVM Phabricator instance.

[mlir][Linalg] Introduce a BufferizationAliasInfo (6/n)
ClosedPublic

Authored by nicolasvasilache on Jun 11 2021, 4:51 AM.

Details

Summary

This revision adds a BufferizationAliasInfo which maintains and updates information about which tensors will alias once bufferized, which bufferized tensors are equivalent to others and how to handle clobbers.

Bufferization greedily tries to bufferize inplace by:

  1. first trying to bufferize SubTensorInsertOp inplace, in reverse order (these are deemed the most expensives).
  2. then trying to bufferize all non SubTensorOp / SubTensorInsertOp, in reverse order.
  3. lastly trying to bufferize all SubTensorOp in reverse order.

Reverse order is a heuristic that seems to work nicely because structured tensor codegen very often proceeds by:

  1. take a subset of a tensor
  2. compute on that subset
  3. insert the result subset into the full tensor and yield a new tensor.

BufferizationAliasInfo + equivalence sets + clobber analysis allows bufferizing nested
subtensor/compute/subtensor_insert sequences inplace to a certain extent.
To fully realize inplace bufferization, additional container-containee analysis will be necessary and is left for a subsequent commit.

Diff Detail

Event Timeline

nicolasvasilache requested review of this revision.Jun 11 2021, 4:51 AM

Drop spurious test file.

This still requires a bit more work and tests but should already be reviewable.

Thanks! Starting working through, mostly stylistic nits while digesting the design.

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
133

Does the printer for BlockArgument add a newline? I don't think it should if it does.

164

Was this supposed to be debug logging?

364–367
379

nit: Drop the llvm:: here.

(This also applies quite a bit in this revision, many of these classes are re-exported in the mlir:: namespace).

384
417

Can you spell out auto here?

431
443

Same here.

456–457

Same comment here as above.

822

There is a mismatch of commenting style in this file, a lot of //.

843

Why the inline here?

864–865

Unfinished comment?

1497–1498

nit: Spell out auto here.

1517–1520

Should this be removed? Looks like a TODO.

1550–1559
nicolasvasilache edited the summary of this revision. (Show Details)Jun 16 2021, 3:57 AM

Finish the implementation and make tests green again.

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
133
nicolasvasilache marked 15 inline comments as done.Jun 16 2021, 4:39 AM
nicolasvasilache added inline comments.
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
456–457

Can't seem to match this comment to a previous one, the code was cleaned up a bunch in the meantime, tentatively marking as resolved.

822

Made a general cleanup of /// for methods and classes in the meantime, I think this comment was referring to those?
Tentatively marking as done.

nicolasvasilache marked 2 inline comments as done.

Address River's comments.

Avoid DenseMap update-while-iterating, which hopfully will fix Windows.
Add constness to relevant places in the analysis following-up on update-while-iterating issue.
Fix one test in comprehensive-func-bufferize which did not consider ReturnOp to read and resulted in a spurious clobber.

ftynse accepted this revision.Jun 17 2021, 2:55 AM
ftynse added inline comments.
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
168

Could we store the result of lookupOrNull in a variable and avoid double lookup here?

297

ReturnOp has no results so it will also be caught by the second alternative below.

363–364

This looks unnecessary and can be treated as "other ops" in the default clause.

487

Nit: explicit

490–491

I couldn't match the documentation with the function name. The documentation says "would bufferize to a write to a buffer", but the function is named "aliasesNonWritableBuffer". How can it be non-writable if there is a write to it?

493–494

Looks like copy-pasta form the above.

497
1013

Leftover debug

1175
1532

Looks like leftover (excessive) debug

1554

Most places in this file use just TensorType but this, and a couple of others, a more specific RankedTensorType. Is this intentional?

1633–1636

Nit: return mlir::failure(result.wasInterrupted())

This revision is now accepted and ready to land.Jun 17 2021, 2:55 AM
nicolasvasilache marked 12 inline comments as done.

Address Alex's comments.

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
497

I thought the function/method comment style guide was to use infinitive rather than 3rd person ?
I've tried to be consistent with use of infinitive.

1554

Yes, there are some differences that will appear more clearly with unranked tensors (e.g. calls to external functions are supported but alloc of unranked buffer for out-of-place unranked tensor is not).
Good parts of this code are ported from an external prototype where unranked tensors are supported to some extent; I'll be able to tighten all uses when I get to adding the right functionality
In this particular case it seems TensorType is more permissive and is the right call.

Revert spuriously commented line.

ftynse added inline comments.Jun 18 2021, 12:31 AM
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
497

This was (and still is) saying "Merge <...> and iterates" <...>", which is grammatically incorrect regardless of your preference for the comment style. I don't care beyond local consistency.

nicolasvasilache marked an inline comment as done.

Pass bufferAliasInfo to bufferize functions to fix subtensor_insert bufferization impl. that does not match the description.

This revision was landed with ongoing or failed builds.Jun 21 2021, 12:04 AM
This revision was automatically updated to reflect the committed changes.
mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferize.cpp
497

Fixed.