iree icon indicating copy to clipboard operation
iree copied to clipboard

[Stream] Update dispatch sites encodings in UnifyEncodingForGlobals pass 2/n

Open hanhanW opened this issue 2 weeks ago • 1 comments

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

hanhanW avatar Dec 11 '25 08:12 hanhanW

I found few issues when I added few more tests. I'll open it after I fix it.

hanhanW avatar Dec 11 '25 12:12 hanhanW