Add lowerings for mma, register and allocate #86
Conversation
This PR adds a mma unit test which lowers to vector.loads/stores and amdgpu.mfmas. Also supports shared memory promotion. Signed-off-by: Harsh Menon <harsh@nod-labs.com>
Signed-off-by: Harsh Menon <harsh@nod-labs.com>
| emitter.emit(graph.get_root_graph()) | ||
| emitter.finish() | ||
|
|
||
| if kwargs.get("canonicalize", False): |
There was a problem hiding this comment.
When do we not want to canonicalize?
There was a problem hiding this comment.
You don't want to canonicalize when you will lose all your IR on canonicalization (because they have no uses). This happens on some of the other tests (you can try it out by setting canonicalize=True on some of the other tests).
| named_sequence = transform_d.NamedSequenceOp( | ||
| "__transform_main", [any_op_t()], [] | ||
| ) | ||
| with InsertionPoint(named_sequence.body): |
There was a problem hiding this comment.
Just curious, is this the best way to setup canonicalization? by setting up a TD-like structure to do the canonicalization? I'd assume we'd have better API support from upstream to call canonicalization and or other standard patterns on.
There was a problem hiding this comment.
Good question and I was thinking about this myself. I think we can avoid TD by using other patterns, but there could be some advantages to using TD for now (as it has support for other patterns like LICM etc.). If we discover that we don't need that, we can always rewrite this without using TD.
| raise CodegenError("No hardware constraints found.") | ||
|
|
||
| result = None | ||
| for constraint in hardware_constraints: |
There was a problem hiding this comment.
We probably should validate len(hardware_constraints) == 1 and get rid of this loop.
There was a problem hiding this comment.
Sure, will add this to the follow on PR.
This PR adds a mma unit test which lowers to
vector.loads/stores and amdgpu.mfmas. Also supports shared memory promotion.