Add DeviceTensorTrait to export annotations

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

iree/turbine/aot/support/ir_utils.py

@@ -67,6 +67,7 @@
 
 from ..tensor_traits import (
     DeviceAffinity,
+    DeviceTensorTrait,
     ExternalTensorTrait,
 )
 
@@ -272,6 +273,7 @@ def create_tensor_global(
     ) -> Tuple[str, Operation, IrType]:
         element_type = self.torch_dtype_to_iree_type(t.dtype)
         external, external_scope, external_name = attrs.infer_external_from_tensor(t)
+        device = DeviceTensorTrait.get(t)
 
         # Always create globals at the top. Then after created, if there was
         # a prior one, move the new one to after it to maintain declaration
@@ -287,6 +289,10 @@ def create_tensor_global(
                 ir_attrs["noinline"] = UnitAttr.get()
             if attrs.mutable:
                 ir_attrs["is_mutable"] = UnitAttr.get()
+            if device:
+                ir_attrs["iree.abi.affinity"] = Attribute.parse(
+                    f"#hal.device.promise<@{device.device_id}>"
+                )
             if external:
                 # Emit named external reference.
                 external_scope_attr = StringAttr.get(external_scope or "model")

iree/turbine/aot/tensor_traits.py

@@ -12,6 +12,7 @@
 
 __all__ = [
     "DeviceAffinity",
+    "DeviceTensorTrait",
     "ExternalTensorTrait",
 ]
 
@@ -31,6 +32,26 @@ def __repr__(self) -> str:
         return f"DeviceAffinity({self.ordinal})"
 
 
+@dataclass
+class DeviceTensorTrait:
+    """Represents a 'trait' that can be applied to a Tensor to signal that
+    it is to be loaded to a speific device at execution time.
+    """
+
+    device_id: str
+
+    @staticmethod
+    def get(from_tensor: torch.Tensor) -> Optional["DeviceTensorTrait"]:
+        existing = getattr(from_tensor, "_turbine_device_tensor_trait", None)
+        if existing is None:
+            return None
+        assert isinstance(existing, DeviceTensorTrait)
+        return existing
+
+    def set(self, to_tensor: torch.Tensor):
+        to_tensor._turbine_device_tensor_trait = self  # type: ignore
+
+
 @dataclass
 class ExternalTensorTrait:
     """Represents a 'trait' that can be applied to a Tensor to signal that

tests/aot/params_test.py

@@ -16,6 +16,7 @@
     export,
     externalize_module_parameters,
     save_module_parameters,
+    DeviceTensorTrait,
     ExternalTensorTrait,
     ParameterArchive,
     ParameterArchiveBuilder,
@@ -119,6 +120,15 @@ def testExternalTensorTrait(self):
         self.assertIs(ExternalTensorTrait.get(t), trait)
 
 
+class DeviceTensorTest(unittest.TestCase):
+    def testDeviceTensorTrait(self):
+        t = torch.ones([2, 3], dtype=torch.float32)
+        trait = DeviceTensorTrait(device_id="__device_who")
+        self.assertIsNone(trait.get(t))
+        trait.set(t)
+        self.assertIs(DeviceTensorTrait.get(t), trait)
+
+
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
     unittest.main()