tree f4ecb9f287801510dfe34c522136606f7ae29f7a
parent 01d3cd5c9624bf721f4e86a0dca3da897c76bdd8
author ziereis <44057120+ziereis@users.noreply.github.com> 1785392730 +0200
committer GitHub <noreply@github.com> 1785392730 +0200
gpgsig -----BEGIN PGP SIGNATURE-----
 
 wsFcBAABCAAQBQJqau5aCRC1aQ7uu5UhlAAAXvsQACNkXWxTW4jL3PMWLveFKgH5
 3w9PvMUMCIv75bkgpwNM5x0znNRHl3MBdKXgcoztfpOmcZmDtt728JgQcNWgr2Jm
 SsO7APqpI4dcexW6zVAf6zVJv/74bUqqLYFF0Eg0T1L7/W8S+nn045yhOB2/Ihoe
 IuSSozvaWckUXlPomWoZYS8/fbG+y0qV+QmfBfxZAX3n3JM180kk8DtRW7Ouc/c1
 Axd38OOkHDQJEFFKhYFSdVgdPbJ3Tv6p2b5fgf7snwi572MajR1VEQAkH6zWuswZ
 ijhTltj4KRbIoX3JjryS51DYm/Msk09zgYMBXoAKGn+jALzohbMkaBQPjCPzhVkF
 Oxns1IJlqyVomoEzoqVEAItXpnvWasYeQP4Rfskt/r87ARFhNVMbmTXjxY6MSd6p
 DF775trDdEBFymX6u84phxK48Ka0R4uVrq4UjkYj/b5nxRQ/iB/BcuguuJppmtoG
 zh/8FY+5QwoEZIyZ+Vtwag+gPnINl3RNxBsGgtGC84ma7gzZU6d7ZH4bLkqLLHjg
 QV3LuGEo4AhfTZWJW5APusZbthoep6AtS8I+i4YB6buznhnupBejOoR6BaQyBbqG
 ADO+HjJAtznWrhly2+ce5q8Lk7ioJ7b0iz144pJnxzUK6YVU64nw5bmT5Ak+hFGr
 giel42nRi9tmDP8cOp2m
 =AKu6
 -----END PGP SIGNATURE-----
 

 [Stream] Model execution affinity for cross-device transfers (#24748)

Fixes #21361.

A `stream.async.transfer` currently carries source and target
affinities, but those describe where its input and result reside, not
necessarily where the transfer executes.

This distinction matters when a value produced on device B is
transferred into a result pinned to device A. The problems this can
cause are described in detail in the issue linked above.

This change adds an optional execution affinity to
`stream.async.transfer`, printed as `on(...)`, alongside the existing
`from(...)` and `to(...)` affinities:

  ```mlir
  %result = stream.async.transfer
      on(#hal.device.affinity<@device_a>)
      %source
      from(#hal.device.affinity<@device_b>)
      -> to(#hal.device.affinity<@device_a>)
      !stream.resource<external>
```

A new PlaceTransferExecutionsPass runs before execution scheduling. It uses affinity analysis to identify transfers whose result has a unique pinned affinity and places those transfers on the pinned device when the default producer-side placement is incompatible.

ScheduleExecutionPass then handles the transfer through its normal affinity partitioning. Transfers are only grouped with their producer when their execution affinities are compatible.

The existing staging-transfer placement rules are preserved as the default when no explicit execution affinity is present. Source and target affinities are no longer modified indirectly when setting an operation's execution affinity.

The full Stream pipeline now produces the desired result for the original multi-device case:

  - The device-B dispatch result is allocated with optimal<A, B> affinity.
  - The final external result remains pinned to device A.
  - The transfer into that result executes on device A.
  
  With this https://github.com/iree-org/iree/pull/21589 can also be closed as this implementation should solve the problem.

---------

Signed-off-by: default <ziereis@roofline.ai>