Page MenuHomePhabricator

[mlir][Affine] Introduce affine.vector_load and affine.vector_store
ClosedPublic

Authored by dcaballe on Fri, May 8, 6:22 PM.

Details

Summary

This patch is the first of a short series that aims to extend Affine dialect
with vector memory semantics. Further context and details can be found here:
https://llvm.discourse.group/t/understanding-the-vector-abstraction-in-mlir/

In particular, this patch adds affine.vector_load and affine.vector_store ops to
the Affine dialect and lowers them to vector.transfer_read and
vector.transfer_write, respectively, in the Vector dialect.

Subsequent patches will add new interfaces for affine load-like and
store-like ops (already implemented locally), port affine fusion to
these interfaces as an example (already implemented locally), and
extend affine analyses to deal with vector semantics (WIP, may need
some help :)). Some refactoring may also be needed.

Thanks!
Diego

Diff Detail

Event Timeline

There are a very large number of changes, so older changes are hidden. Show Older Changes
dcaballe added inline comments.Fri, May 8, 6:43 PM
mlir/include/mlir/Conversion/Passes.td
18 ↗(On Diff #262985)

Should we rename this to convert-affine-to-loops?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
811

Not sure if it's feasible at all in tablegen but it would be great if we could make "code" variables (extraClassDeclaration, builders, etc.) "additive" in a hierarchy instead of overriding the variable with the value from the subclass.

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp
55 ↗(On Diff #262985)

I would be great if we could have a builder that sets getMinorIdentityMap as the default value of the permutation map for those cases where no permutation is needed.

Thanks for initiating this - this would be a great addition to the affine dialect.

mlir/include/mlir/Conversion/Passes.td
18 ↗(On Diff #262985)

The reasons this is still called lower-affine is that it's actually converting affine not just to SCF (formerly loop) but SCF + std. Because the affine.apply's as well as the expressions in the affine for bounds, affine if conditionals, and load/store subscripts are lowered to add/mul/div etc. -convert-affine-to-<something> is perhaps more consistent / better.

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp
55 ↗(On Diff #262985)

+1.

bondhugula requested changes to this revision.Fri, May 8, 7:00 PM
bondhugula added inline comments.
mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h
1 ↗(On Diff #262985)

Instead of another conversion pass for this, we could just do this as part of -lower-affine? Because lower-affine already converts to scf + std, and vector is already at the same "level" of abstraction as std, so, -lower-affine could as well just convert to scf + std + vector.

mlir/test/Dialect/Affine/load-store.mlir
228

Do you want to switch the type order in the type list to keep it consistent with affine.load? Just makes it easier to copy / paste while writing test cases and a bit consistent. The order of the operands need not be the same as in the colon type list, but it's the semantics "store to a memref of type ..., a vector value of type ..." that we'd read.

This revision now requires changes to proceed.Fri, May 8, 7:00 PM
dcaballe marked 3 inline comments as done.Mon, May 11, 11:15 AM
dcaballe added inline comments.
mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h
1 ↗(On Diff #262985)

Ok! I did it this way for composability reasons since I need the vector part but not necessarily the scf + std. I'll move it to -lower-affine but leave the populateAffineToVectorConversionPatterns around so that I can create my own lowering pass locally.

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp
55 ↗(On Diff #262985)

Ok, I'll address this separately

mlir/test/Dialect/Affine/load-store.mlir
228

Not sure about this. I tried to keep it consistent with vector.transfer_write so that we don't have to do a different interpretation of a vector store depending on the dialect we are. It could be confusing if we switch the type order since one would initially expect the types to match the order of the operands when the number of operands and types are the same (at least I would :)). The way I read it is more like: "store a vector value of type ... to a memref of type...".

Great, thanks @dcaballe !

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
811

In base:

let extraClassDeclarationBase = " ... "

In derived:

let extraClassDeclaration =  extraClassDeclarationBase # "..."

?

mlir/lib/Conversion/AffineToVector/AffineToVector.cpp
55 ↗(On Diff #262985)

Feel free to also fold it into this revision if you want.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2502

Not for this revision but generally, should printer/parser based on parseAffineMapOfSSAIds have a special syntax in the declarative assembly, e.g. "$affine_indexing" ?

mlir/test/Dialect/Affine/load-store.mlir
228

feel free to change this syntax, parser and printer + the vector_transfer syntax to whichever combination of you prefer:
vector from memref / vector to memref or memref to vector / memref from vector

dcaballe marked an inline comment as done.Mon, May 11, 1:37 PM
dcaballe added inline comments.
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
811

Something like that should work! However, tablegen complains about extraClassDeclarationBase not being defined. I can file a bugzilla for this. I wouldn't like to go into extending ODS in this patch. I could keep getVectorType in the base class for now and point to bugzilla.

dcaballe updated this revision to Diff 263306.Mon, May 11, 4:21 PM
dcaballe marked 2 inline comments as done.

Addressing feedback. Thanks!

dcaballe added inline comments.Mon, May 11, 4:23 PM
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
811

Also note that part of the base extraClassDeclaration, if not all, will go to the interfaces for affine loads/stores (next patch).

mlir/test/Dialect/Affine/load-store.mlir
228

No strong opinion here but the current combination in vector.transfer_* seems ok to me. In any case, if we want to change it, probably better to do it in a separate patch since we would have to change vector tests, etc.

Just a nit but the one letter difference between affine.vstore and affine.store does not seem ideal to me from a readability point of view. Also one-letter abbreviation isn't great in general, "v" could be for "volatile" for example.
What about affine.vector_store (or affine.vec_store)?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
803

Typo: affine.load -> affine.vload

817

Do we have another term than "blocking"? It can be ambiguous with respect "synchronous" for example, and it isn't clear to me in general.
Is this "store a vector to a contiguous memory segment"?

mlir/include/mlir/Dialect/Vector/VectorOps.td
998

Typo permunation

1069

Typo permunation

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2502

I thought the same! Can you file a bug for this?

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir
9

Isn't this reading out-of-bound?

dcaballe marked 2 inline comments as done.Tue, May 12, 9:00 AM

Thanks, Mehdi!

What about affine.vector_store (or affine.vec_store)?

affine.vector_store sounds good to me. Anybody else?

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
817

This doc is a copy-paste mix of AffineStoreOp doc and VectorTransferWriteOp doc, including the "blocking" terminology. I decided to do that for consistent with the vector dialect. I can replace "blocking write" by something along the lines of "store to a vector to a contiguous memory segment" in both dialects if that is clearer. We may have to revisit this doc anyways if/when we introduce strided or gather/scatter accesses.

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir
9

It might anyways due to arg0 being an arbitrary input but I'll make it less obvious :)

bondhugula added inline comments.Tue, May 12, 9:13 AM
mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h
1 ↗(On Diff #262985)

Yes, that sounds good. You don't need a separate pass. Curious: why do you need just the conversion to vector read/write transfer. What will the output give you that you won't get by looking at affine.vload/vstore?

mlir/test/Dialect/Affine/load-store.mlir
228

IMHO, it's more consistent and readable to always find the memref first and the vector type trailing it. The order of operands is secondary - for example, the load doesn't even have a vector type operand. The type list is all custom in the custom form, and so consistency for all memory ops would be good.

bondhugula requested changes to this revision.Tue, May 12, 9:14 AM

Please update the commit summary that there is no new conversion pass.

This revision now requires changes to proceed.Tue, May 12, 9:14 AM
ftynse added a subscriber: ftynse.Tue, May 12, 10:02 AM
ftynse added inline comments.
mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
809

This can be omitted.

811

However, tablegen complains about extraClassDeclarationBase not being defined.

let _redefines_ a variable, <typename> defines a new one. Try code extraClassDeclarationBase = [{...}]; instead.

817

I think the "blocking" terminology in vector transfers indicates that it's not a DMA operation and it will block execution until the result is ready. It doesn't make sense for load/store.

mlir/include/mlir/Dialect/Vector/VectorOps.td
1002

This looks complex enough to deserve being in a .cpp file instead.

dcaballe edited the summary of this revision. (Show Details)Tue, May 12, 11:20 AM
dcaballe updated this revision to Diff 263508.Tue, May 12, 1:30 PM
dcaballe marked 10 inline comments as done.

Addressing all the feedback. Thanks!

mlir/include/mlir/Conversion/AffineToVector/AffineToVector.h
1 ↗(On Diff #262985)

I meant that I may want to compose the affine->vector lowering with the lowering to other dialects or a customized scf+std lowering.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
811

Awesome! It works, thanks!

817

I removed "blocking" and rephrase it a bit, describing the read/write starting position and the shape of the slice being read/written. Please, let me know if you have more comments. We can iterate on this.

mlir/test/Dialect/Affine/load-store.mlir
228

Ok, I changed that for affine.vector_load and affine.vector_store. We can change vector.transfer_* separately.

dcaballe marked an inline comment as done.Tue, May 12, 3:01 PM
dcaballe added inline comments.
mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2502

Thanks, LGTM but probably worth having Uday or Alex approval on this.

Also: you did mention in the past that you were interested to look into gather/scatter load/store with some 2D native memory unit in mind, do you still plan to look into this later?

Thanks, Mehdi!

Also: you did mention in the past that you were interested to look into gather/scatter load/store with some 2D native memory unit in mind, do you still plan to look into this later?

The 2D native memory unit was just an example. I'm interested in having a more generic vector representation, including strided loads/stores, generic gathers/scatters, and multi-dimensional vector loads/stores in general. However, it depends on priorities and my workload. I think there are a few things that I may need to look into first around the current affine vector ops before going any further. I'll be happy to help, though, if somebody else wants to champion this.

bondhugula accepted this revision.Tue, May 12, 11:30 PM

This overall looks great. Please also update the commit title for vload/vstore -> vector_load/store. (since arc won't do it by itself as you know.)

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
783

Nit: I assume you also want strides in the future. You could document somewhere that this slice/vector is currently contiguous along the respective dimensions, but will be extended to strided vectors.

826

Nit: "... affine function of ... " or "... affine expression in ..."

827

Nit: ... the start position of the write within ...

mlir/include/mlir/Dialect/Vector/VectorOps.td
1002

Was a build method moved into AffineOps.cpp? I don't see anything there.

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
579

Argument comment for 1 (i.e., /*numResults=*/1). I think we should also improve the doc comment for getMinorIdentityMap. It doesn't say what 'results' is, and it should be numResults.

607

Comment update: std.store ->

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2572

We are missing a check to verify that this op's map is indeed a permutation map. Is vector.read/write_transfer also missing it? If yes, please add a TODO for it on AffineVector/Load/Store Op's verifiers. In future, I assume you plan to support strides as well - so that could be marked as a TODO as well.

mlir/test/Conversion/AffineToStandard/lower-affine-to-vector.mlir
9

-memref-bound-check can catch such out-of-bound errors (but for affine.load/store ops).

mlir/test/Dialect/Affine/load-store.mlir
220

Nit: MAP0 -> MAP_ID or IDENTITY for readability.

257

We could add an extra negative test case to test/Dialect/Affine/invalid.mlir where you use a different element type for the vector and for the memref, and check for errors.
affine.vector_load %A[%i] : memref<100xvector<8xf32>>, vector<16xf32>
affine.vector_store %v, %A : memref<100xf32>, vector<8xf64>
affine.vector_load %A[%i] : memref<100xvector<8xf32>>, vector<8xf32>

All of these should fail verification. Very likely that users will try these. :)

This revision is now accepted and ready to land.Tue, May 12, 11:30 PM
dcaballe retitled this revision from [mlir][Affine] Introduce affine.vload and affine.vstore to [mlir][Affine] Introduce affine.vector_load and affine.vector_store.Wed, May 13, 9:32 AM
dcaballe updated this revision to Diff 263888.Wed, May 13, 4:44 PM
dcaballe marked 9 inline comments as done.

Addressing feedback.

  • Added specific verifiers for affine.vector_load and affine.vector_store.
  • Refactored affine.load and affine.store verifiers to avoid code duplication.
  • Added invalid tests.

This overall looks great. Please also update the commit title for vload/vstore -> vector_load/store. (since arc won't do it by itself as you know.)

Thanks! Yes, I'll clean up the message before committing to remove irrelevant information.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
826

I think 'expression' is referring to the AffineExpr class here (note that this is copy-pasted from AffineStoreOp). "affine expression _in_ loop induction variables and symbols" sounds odd to me but I can be wrong. Is that what you want?

mlir/include/mlir/Dialect/Vector/VectorOps.td
1002

It was moved to VectorOps.cpp since this is VectorOps.td

mlir/lib/Conversion/AffineToStandard/AffineToStandard.cpp
579

We can actually remove all of this and use the new build with default permutation and padding.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2572

The current map in affine.vector_load/store is not a permutation map (at least a permutation map in the way is defined in vector.transfer_load/store), right? IIUC, the one in affine.vector_load/store is the same as in affine.load/store, i.e., it's just for indexing (computing the start position of the read/write), not for permuting the slice we are reading/writing from/to memory. I think we would need a separate map for that in affine, if we are interested in modeling that.

Added TODOs for both things to the doc in the td file but, please, let me know if I'm wrong.

mlir/test/Dialect/Affine/load-store.mlir
257

Done. I added specific verifiers since this was not covered by the auto-generated ones. This required some refactoring of the scalar verifiers since I didn't want to have the same code replicated 4 times.

bondhugula accepted this revision.Wed, May 13, 8:34 PM
bondhugula marked 2 inline comments as done.

Looks good. All your examples/test cases are with 1-d vectors, but the op documentation isn't restricted. You may just want to either specifically add that the vector could be multi-dimensional or include one test case in load-store.mlir (or example in the op doc). I assume such multi-dimensional affine vector load/stores correctly lower to vector read/write transfer as is with what you have.

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
826

Never mind - this is too minor. Sounds good as is.

mlir/lib/Dialect/Affine/IR/AffineOps.cpp
2572

Oh, I see. Looks like I got confused. There is no permutation map here.

dcaballe updated this revision to Diff 263922.Wed, May 13, 11:18 PM

Enabled multi-dim vector loads/stores (minor fixes)
Added tests and doc examples.

I'll commit it tomorrow if no more comments.
Thanks!

This revision was automatically updated to reflect the committed changes.