dialects: (affine) Add the VectorLoadOp and VectorStoreOp operations#5365
dialects: (affine) Add the VectorLoadOp and VectorStoreOp operations#5365jakedves wants to merge 4 commits intoxdslproject:mainfrom
VectorLoadOp and VectorStoreOp operations#5365Conversation
Co-Authored-By: Nick Brown <11333808+mesham@users.noreply.github.com>
Co-Authored-By: Nick Brown <11333808+mesham@users.noreply.github.com>
Co-Authored-By: Nick Brown <11333808+mesham@users.noreply.github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #5365 +/- ##
==========================================
- Coverage 85.91% 85.90% -0.02%
==========================================
Files 378 378
Lines 53986 54016 +30
Branches 6178 6181 +3
==========================================
+ Hits 46384 46401 +17
- Misses 6127 6140 +13
Partials 1475 1475 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| func.func @affine_vector_load() -> vector<8xf32> { | ||
| %0 = "test.op"() : () -> memref<100x100xf32> | ||
| // CHECK: %{{.*}} = "test.op"() : () -> memref<100x100xf32> | ||
| %1 = "affine.vector_load"(%0) <{map = affine_map<() -> (1, 2)>}> : (memref<100x100xf32>) -> vector<8xf32> |
There was a problem hiding this comment.
can you please add custom syntax for these?
There was a problem hiding this comment.
This is quite difficult to handle with the syntax so may have to come back to this later.
I think this would be simple with a way to parse the affine expressions which I think exists, but can't find the exact call to use.
Some reference test cases that work in MLIR:
%0 = memref.alloc() : memref<100x100xf32>
affine.for %i0 = 0 to 10 {
affine.for %i1 = 0 to 16 step 4 {
%1 = affine.vector_load %0[3, 7] : memref<100x100xf32>, vector<4xf32>
%2 = affine.vector_load %0[%i0, %i1] : memref<100x100xf32>, vector<4xf32>
%3 = affine.vector_load %0[%i0, 7] : memref<100x100xf32>, vector<4xf32>
%4 = affine.vector_load %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
%5 = affine.vector_load %0[%i0, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
%6 = affine.vector_load %0[3 + %i1, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
%7 = affine.vector_load %0[3 + %i0 * 7 + %i1, 7] : memref<100x100xf32>, vector<4xf32>
affine.vector_store %1, %0[%i0 + 3, %i1 + 7] : memref<100x100xf32>, vector<4xf32>
}
}There was a problem hiding this comment.
I can sort of imagine how to do it, both in a hacky way and an unhacky way
There was a problem hiding this comment.
The hacky way would be to do string substitution to map s0 to %i0 and so on
There was a problem hiding this comment.
I just realized that our existing implementation of vector.transfer_read and vector.transfer_write, which are also supposed to support this system, are broken... It would be great to handle this properly. If you have the time to have a stab at this in a separate PR fixing the custom syntax in these ops that would be great, otherwise I'll try to have a go this weekend.
There was a problem hiding this comment.
I could probably have a go but I'm not free for this until sometime next week
There was a problem hiding this comment.
I finally have time, will look at it over the next week
There was a problem hiding this comment.
Hi @jakedves! Happy new year :) Do you have some time to look at this in the next few weeks?
| memref: SSAValue, | ||
| indices: Sequence[SSAValue], | ||
| map: AffineMapAttr | None = None, | ||
| result_type: Attribute | None = None, | ||
| ): | ||
| if not isa(memref_type := memref.type, MemRefType[TypeAttribute]): | ||
| raise ValueError( | ||
| "affine.vector_load memref operand must be of type MemRefType" | ||
| ) |
There was a problem hiding this comment.
| memref: SSAValue, | |
| indices: Sequence[SSAValue], | |
| map: AffineMapAttr | None = None, | |
| result_type: Attribute | None = None, | |
| ): | |
| if not isa(memref_type := memref.type, MemRefType[TypeAttribute]): | |
| raise ValueError( | |
| "affine.vector_load memref operand must be of type MemRefType" | |
| ) | |
| memref: SSAValue[MemRefType], | |
| indices: Sequence[SSAValue], | |
| map: AffineMapAttr | None = None, | |
| result_type: Attribute | None = None, | |
| ): |
Possibly cleaner?
| if map is None: | ||
| # Create identity map for memrefs with at least one dimension or () -> () | ||
| # for zero-dimensional memrefs. | ||
| if not isinstance(memref_type := memref.type, MemRefType): |
There was a problem hiding this comment.
Feels like the same could be done here
3 participants
Adds the vector load/store ops to the affine dialect as in MLIR, with xDSL and MLIR roundtrip tests