This is an archive of the discontinued LLVM Phabricator instance.

[mlir][MemRef] Add patterns to extract address computations
ClosedPublic

Authored by qcolombet on Mar 23 2023, 7:52 AM.

Details

Summary

This patch adds patterns to rewrite memory accesses such that the resulting
accesses are only using a base pointer.
E.g.,

mlir
memref.load %base[%off0, ...]

Will be rewritten in:

mlir
%new_base = memref.subview %base[%off0,...][1,...][1,...]
memref.load %new_base[%c0,...]

The idea behind these patterns is to offer a way to more gradually lower
address computations.

These patterns are the exact opposite of FoldMemRefAliasOps.
I've implemented the support of only three operations in this patch:

  • memref.load
  • memref.store
  • nvgpu.ldmatrix

Going forward we may want to provide an interface for these rewritings (and
the ones in FoldMemRefAliasOps.)
One step at a time!

Diff Detail

Event Timeline

qcolombet created this revision.Mar 23 2023, 7:52 AM
qcolombet requested review of this revision.Mar 23 2023, 7:52 AM
nicolasvasilache accepted this revision.Mar 23 2023, 8:02 AM

Thanks, LGTM, make sense to keep the interface work for later.

One note though, we also need those on vector.transfer soonish.

mlir/include/mlir/Dialect/MemRef/Transforms/Passes.h
65 ↗(On Diff #507740)

Similary to https://reviews.llvm.org/D146624 we want all these populate to live in Transform/Transforms.h

mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
7

It would be nice to find a name that somehow ties these logically to FoldMemRefAliasOps.cpp but I odn't have a better suggestion atm.

This revision is now accepted and ready to land.Mar 23 2023, 8:02 AM
qcolombet updated this revision to Diff 508029.Mar 24 2023, 3:47 AM
  • Add support for vector.transfer_read/write
  • Move the populate function in its own header (we'll want to move others as well in an NFC patch)
  • Fix the dropping of the nontemporal flag (and add a test for that)
  • Refactor getSrcMemRef to not choke on tensor type
  • Refactor getViewSizeForEachDim to use the same implementation for ldmatrix, vector.transfer_read/write
  • Add tests for vector.transfer_read/write
qcolombet updated this revision to Diff 508034.Mar 24 2023, 3:52 AM
  • Add missing vector dialect dependency to the test-extract-address-computations pass
qcolombet marked an inline comment as done.Mar 24 2023, 3:59 AM
qcolombet added inline comments.
mlir/lib/Dialect/MemRef/Transforms/CMakeLists.txt
7

Yeah, I can't find anything else either.
Maybe we could rename the other one?

No good idea there either though!

nicolasvasilache accepted this revision.Mar 24 2023, 4:24 AM

thanks for adding vector.transfer too, LGTM!

nicolasvasilache requested changes to this revision.Mar 24 2023, 4:33 AM

sorry clicked too soon, added some comments

mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
36

also move the other populates here ?

mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
119

For my own education, could you please point me to explanations about \see and other similar things I have seen you use over time ?
I'd happily follow suit and try to improve my own comments.

224

llvm::function_ref instad of C style funptrs ?

247

return rewriter.notifyMatchFailure(op, "...") everywhere plz, this greatly simplifies understanding when debugging

270

Add a good comment that strides are handled properly by composition of subviews here please.
This is not true for the reverse transformation and the precondition for vector.transfer checks and fails to match if strides are not all one: the forward and inverse problems are not symmetrical unfortunately.

mlir/test/Dialect/MemRef/extract-address-computations.mlir
288

some examples with strides please that mirror my previous comment

mlir/test/lib/Dialect/MemRef/TestExtractAddressComputations.cpp
30

can we just wrap this in a transform op and avoid creating unused code?

This revision now requires changes to proceed.Mar 24 2023, 4:33 AM
qcolombet added inline comments.Mar 24 2023, 6:20 AM
mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
36

I thought we would do that in an NFC commit.
Which I'll do after that one land.

mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
119

Here is the complete doc for the doxygen "keywords":
https://www.doxygen.nl/manual/commands.html

The ones I use that may not be that common are:

  • \pre, \post: Document a pre-post condition
  • \see: Create a link to another item (class, function, etc.) also documented in doxygen

But maybe what I think is common or not is biased :P (\note, \code, \p, etc.).

224

Good point, I always forgot about these!
(I start with std::function, then remember that templates don't like them and resort to plain C. I need to update my mental model :)).

247

Good point!

270

👍

mlir/test/Dialect/MemRef/extract-address-computations.mlir
288

👍

mlir/test/lib/Dialect/MemRef/TestExtractAddressComputations.cpp
30

Sure thing.

qcolombet updated this revision to Diff 508123.Mar 24 2023, 9:10 AM
  • Remove the test pass
  • Add a tranform op instead
  • Add a comment for strides when creating the subview
  • Add a test for non-1 strides (only for vector.transfer_write, @nicolasvasilache let me know if you want more)
  • Report failures with notifyMatchFailure
qcolombet marked 3 inline comments as done.Mar 24 2023, 9:12 AM
qcolombet added inline comments.
mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
224

Alright looks like we can't use them either. They suffer the same issue as std::function:

error: non-type template parameters of class type only available with ‘-std=c++20’ or ‘-std=gnu++20’

Am I missing something?

mlir/test/Dialect/MemRef/extract-address-computations.mlir
288

Only did one for vector.transfer_write (see test_transfer_write_op_with_strides).
Let me know if you want some for the others.

nicolasvasilache accepted this revision.Mar 24 2023, 9:41 AM
nicolasvasilache added inline comments.
mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
36

sgtm!

mlir/lib/Dialect/MemRef/Transforms/ExtractAddressComputations.cpp
224

oh well ..

mlir/test/Dialect/MemRef/extract-address-computations.mlir
288

that's great, thanks!

This revision is now accepted and ready to land.Mar 24 2023, 9:41 AM
qcolombet updated this revision to Diff 508946.Mar 28 2023, 3:24 AM
  • Fix bazel build
qcolombet added inline comments.Mar 28 2023, 6:24 AM
mlir/include/mlir/Dialect/MemRef/Transforms/Transforms.h
36

Done in faafd26c4d58