[water] Add lowering pattern using LDS promotion for write group #500
Conversation
This attribute contains information needed for picking different lowering patterns for wave.write ops. Signed-off-by: tyb0807 <[email protected]>
Signed-off-by: tyb0807 <[email protected]>
| auto vecType = cast<VectorType>(vec.getType()); | ||
| Value value = adaptor.getValueToStore(); | ||
|
|
||
| // All wave.write operations require index attributes |
There was a problem hiding this comment.
| // All wave.write operations require index attributes | |
| // All wave.write operations require index attributes. |
Nit: let's be consistent about using full sentences in comments.
| /// DESIGN CHOICE: All operations in the same LDS promotion group must be in | ||
| /// the same scope. This is required because: |
There was a problem hiding this comment.
Is this something that can be checked by the verifier? (I'm assuming scope = block?)
| func::FuncOp funcOp = firstOp->getParentOfType<func::FuncOp>(); | ||
| assert(funcOp && "wave.write operation must be within a function"); |
There was a problem hiding this comment.
Do we care about it being a func.func? Can we just take the first IsolatedFromAbove parent instead?
| assert(funcOp && "wave.write operation must be within a function"); | ||
|
|
||
| // Walk the entire function to find all operations in this group. | ||
| funcOp.walk([&](wave::WriteOp writeOp) { |
There was a problem hiding this comment.
Non-local patterns are generally frowned upon...
| if (!ldsAlloc) { | ||
| return rewriter.notifyMatchFailure(firstOp, | ||
| "failed to create LDS allocation"); | ||
| } |
There was a problem hiding this comment.
We can't just notifyMatchFailure here and below because createLdsAllocation created new IR, and there is no rollback possible. Basically, patterns shouldn't return failure after they changed the IR in any way.
|
Generally this looks like too much of an abstraction jump, which will make the code more complex than it should be. IMO, this does not belong in lowering and should happen in two steps: (1) wave-to-wave rewrite that takes a store from registers to global and changes it to store from registers to lds, barrier (needs to be introduced in the wave dialect), load from lds to registers and store from registers to global; then (2) ideally the normal lowering will just work at this point without extra lowering logic. This will also make clear where do all these access patterns come from: they are inferred from the existing load and constraints. Bonus points for a separate PR that introduces an operation that represents lds-to-global copy that (1) can emit. We can then have an option to either expand it into load/store via registers or lower it directly to the appropriate operation when the architecture supports it and it is desirable. |
2 participants
Based on #490.
Served as base for design discussion for now.
DO NOT MERGE.