iree
iree copied to clipboard
[Stream] Update dispatch sites encodings in UnifyEncodingForGlobals pass 2/n
Added functionality to update stream.tensor.dispatch operand encodings when unifying encodings for globals.
- Uses Explorer with walkTransitiveUses for cross-function tracking of encoded global uses.
- Follows stream.async.clone ops (BFS worklist) to collect updates. Currently, only stream.tensor.dispatch op is updated. Other ops that have TensorPhaseOp are not implemented yet.
- Applies all updates in one shot using
DenseMap<Operation*, DenseMap<unsigned, Attribute>>. This makes update more efficient because we don't have to iterate all the encoding operands in every update.
New tests:
- Basic dispatch with encoded operand.
- Multiple encoded operands in same dispatch (@use_both_encodings).
- Cross-function tracking via
util.call(@cross_function_call → @dispatch_helper).
It is a step towards https://github.com/iree-org/iree/issues/22485
I found few issues when I added few more tests. I'll open it after I fix it.