Add `DeviceTensorTrait` to export annotations

iree-org / iree-turbine

IREE's PyTorch Frontend, based on Torch Dynamo.

Apache License 2.0

55 stars 25 forks source link

Add `DeviceTensorTrait` to export annotations #240

Closed rsuderman closed 3 weeks ago

rsuderman commented 3 weeks ago

Exported globals can now be parked to belong to a specific device. This works with sharding tooling to ensure globals are not misassigned to the default affinity device

rsuderman commented 3 weeks ago

We are inconsistent with the DeviceAffinity being an ordinal and this is using a the whole device name directly.

Addressed