Add device affinities for arguments in AOT

We don't have support for providing device affinities for function
arguments, which need to end up as MLIR function argument attributes.

This change adds a class DeviceAffinity and provides the ability to
supply affinities when exporting Torch functions/modules or when
tracing in IREE-Trubine itself.

Signed-off-by: Boian Petkantchin <boian.petkantchin@amd.com>

iree/turbine/aot/compiled_module.py

@@ -44,6 +44,8 @@
     ModuleBuilderOptions,
 )
 
+from .tensor_traits import DeviceAffinity
+
 
 __all__ = [
     "CompiledModule",
@@ -107,12 +109,27 @@ def __call__(self, *args, **kwargs):
         return self.py_value(*args, **kwargs)
 
 
+class ExportTargetDef:
+    def __init__(
+        self,
+        target: Union[Callable, ExportedProgram],
+        *,
+        arg_device: dict[int, DeviceAffinity] | None = None,
+    ):
+        self.target = target
+        self.arg_device = arg_device
+
+    def __call__(self, *args, **kwargs):
+        return self.target(*args, **kwargs)
+
+
 class ExportProcDef:
     __slots__ = [
         "callable",
         "export_name",
         "signature",
         "file_line_loc",
+        "arg_device",
     ]
 
     def __init__(
@@ -122,14 +139,22 @@ def __init__(
         *,
         signature,
         file_line_loc: Optional[Tuple[str, int]] = None,
+        arg_device: dict[int, DeviceAffinity] | None = None,
     ):
         self.export_name = export_name
         self.callable = callable
         self.signature = signature
         self.file_line_loc = file_line_loc
+        self.arg_device = arg_device
 
     def copy(self) -> "ExportProcDef":
-        return ExportProcDef(self.export_name, self.callable, signature=self.signature)
+        return ExportProcDef(
+            self.export_name,
+            self.callable,
+            signature=self.signature,
+            file_line_loc=self.file_line_loc,
+            arg_device=self.arg_device,
+        )
 
     def __repr__(self):
         return f"<def {self.export_name}({self.signature})>"
@@ -142,14 +167,19 @@ def __init__(
         *,
         export_name: Optional[str] = None,
         public: bool = False,
+        arg_device: dict[int, DeviceAffinity] | None = None,
     ):
         self.export_name = export_name
         self.exported_program = ep
         self.public = public
+        self.arg_device = arg_device
 
     def copy(self) -> "ExportedProgramDef":
         return ExportedProgramDef(
-            self.exported_program, export_name=self.export_name, public=self.public
+            self.exported_program,
+            export_name=self.export_name,
+            public=self.public,
+            arg_device=self.arg_device,
         )
 
     def __repr__(self):
@@ -207,6 +237,19 @@ def globals_defs(self) -> Generator[Tuple[str, GlobalsDef], None, None]:
         )  # type: ignore
 
     def def_attribute(self, key, value):
+        if isinstance(value, ExportTargetDef):
+            if not isinstance(value.target, ExportedProgram):
+                # We expect exported function.
+                assert callable(value.target) and inspect.isfunction(value.target)
+                return self.def_export_proc(key, value.target, value.arg_device)
+
+            value = ExportedProgramDef(
+                value.target,
+                export_name=key,
+                public=not key.startswith("_"),
+                arg_device=value.arg_device,
+            )
+
         # Some decorators, the only thing we do is convert them to PyOnlyDef.
         # Do that first so the generic descriptor code below handles them.
         if isinstance(value, builtins.jittable):
@@ -233,6 +276,15 @@ def def_attribute(self, key, value):
             logging.debug("DEFINE PY_ONLY: %s = %r", key, value)
             self.add_export(key, value)
             return value
+        if isinstance(value, ExportTargetDef) and isinstance(
+            value.target, ExportedProgram
+        ):
+            value = ExportedProgramDef(
+                value.target,
+                export_name=key,
+                public=not key.startswith("_"),
+                arg_device=value.arg_device,
+            )
         if isinstance(value, ExportedProgramDef):
             if value.export_name is None:
                 value = value.copy()
@@ -250,7 +302,12 @@ def def_attribute(self, key, value):
             f"compiled module: {value!r}"
         )
 
-    def def_export_proc(self, name, f) -> ExportProcDef:
+    def def_export_proc(
+        self,
+        name,
+        f,
+        arg_device: dict[int, DeviceAffinity] | None = None,
+    ) -> ExportProcDef:
         logging.debug("DEFINE EXPORT: %s = %r", name, f)
         # Get a reasonable location.
         file_line_loc = None
@@ -292,7 +349,13 @@ def def_export_proc(self, name, f) -> ExportProcDef:
                 )
             input_sig.append(param_desc)
 
-        info = ExportProcDef(name, f, signature=input_sig, file_line_loc=file_line_loc)
+        info = ExportProcDef(
+            name,
+            f,
+            signature=input_sig,
+            file_line_loc=file_line_loc,
+            arg_device=arg_device,
+        )
         self.add_export(name, info)
         return info
 
@@ -568,6 +631,20 @@ def save_mlir(inst: "CompiledModule", path: Union[Path, str]):
 
     jittable = staticmethod(builtins.jittable)
 
+    @staticmethod
+    def signature_info(
+        *,
+        arg_device: dict[int, DeviceAffinity] | None = None,
+    ) -> Callable:
+        """Annotate an export target function.
+        This annotation is only required when additional information needs to be
+        provided."""
+
+        def _decorator(f: Callable):
+            return ExportTargetDef(f, arg_device=arg_device)
+
+        return _decorator
+
     def __getattr__(self, name):
         info = CompiledModule.get_info(self)
         try:
@@ -633,6 +710,7 @@ def __new__(
                 ep_def.exported_program,
                 symbol_name=ep_def.export_name or "main",
                 symbol_visibility=None if ep_def.public else "private",
+                arg_device=ep_def.arg_device,
             )
 
         # Instantiate procs.
@@ -661,6 +739,7 @@ def invoke_with_self(*args, **kwargs):
                     posargs=proc_def.signature,
                     kwargs={},  # TODO(#128): kwargs
                     loc=loc,
+                    arg_device=proc_def.arg_device,
                 )
                 trace.trace_py_func(invoke_with_self)
                 info.shadow_dict[key] = _uncallable_public_export

iree/turbine/aot/exporter.py

@@ -32,6 +32,8 @@
 from .fx_programs import FxPrograms
 from . import decompositions
 
+from .tensor_traits import DeviceAffinity
+
 __all__ = [
     "export",
     "ExportOutput",
@@ -177,6 +179,7 @@ def export(
     function_name: Optional[str] = None,
     strict_export: bool = True,
     import_symbolic_shape_expressions: bool = False,
+    arg_device: dict[int, DeviceAffinity] | None = None,
 ) -> ExportOutput:
     """Exports a torch.nn.Module.
 
@@ -199,6 +202,7 @@ def export(
     *,
     module_name: Optional[str] = None,
     function_name: Optional[str] = None,
+    arg_device: dict[int, DeviceAffinity] | None = None,
 ) -> ExportOutput:
     """Exports a single entry-point module consisting of an ExportedProgram."""
     ...
@@ -226,6 +230,7 @@ def export(
     function_name: Optional[str] = None,
     strict_export: bool = True,
     import_symbolic_shape_expressions: bool = False,
+    arg_device: dict[int, DeviceAffinity] | None = None,
 ) -> ExportOutput:
     """Generic export of supported entities.
 
@@ -247,6 +252,10 @@ def export(
         must be empty.
       kwargs: Example keyword arguments.
       dynamic_shapes: Dynamic shape specs to pass to torch.export.
+      arg_device: device affinities for the exported function
+        arguments. On what devices should the program expect its arguments.
+        It is a mapping of argument index to device affinity of the flattened
+        arguments.
 
     Returns:
       An ExportOutput object that wraps the compilation and provides
@@ -266,12 +275,14 @@ def export(
             "This is an experimental feature in PyTorch that the IREE Turbine project is still evaluating. Please report issues or experiences."
         )
 
+    from .compiled_module import ExportTargetDef
+
     TransformedModule: Any
     current_decomps = decompositions.current_aot_decompositions()
     if isinstance(mdl, torch.export.ExportedProgram):
         TransformedModule = CompiledModule.create_from_dict(
             "LambdaCompiledModule",
-            {(function_name or "main"): mdl},
+            {(function_name or "main"): ExportTargetDef(mdl, arg_device=arg_device)},
             export_name=module_name or "module",
             options=ModuleBuilderOptions(
                 import_symbolic_shape_expressions=import_symbolic_shape_expressions,
@@ -311,7 +322,12 @@ def export(
 
         TransformedModule = CompiledModule.create_from_dict(
             "LambdaCompiledModule",
-            {(function_name or "main"): exported_program},
+            {
+                (function_name or "main"): ExportTargetDef(
+                    exported_program,
+                    arg_device=arg_device,
+                )
+            },
             export_name=module_name or "module",
             options=ModuleBuilderOptions(
                 import_symbolic_shape_expressions=import_symbolic_shape_expressions,

iree/turbine/aot/fx_programs.py

@@ -14,13 +14,20 @@
 import os
 from pathlib import Path
 from typing import Any, Optional, Union
+from .compiled_module import ExportTargetDef
 
 import functools
 
 import torch
 import torch.nn as nn
 
 from .decompositions import current_aot_decompositions
+from .tensor_traits import DeviceAffinity
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .compiled_module import ExportTargetDef
 
 # The dynamic_shapes support showed up in the Torch 2.3 timeframe.
 _supports_dynamic_shapes = hasattr(torch.export, "Dim")
@@ -61,7 +68,7 @@ class FxPrograms:
     """
 
     def __init__(self):
-        self.programs: dict[str, torch.export.ExportedProgram] = {}
+        self.programs: dict[str, ExportTargetDef] = {}
 
     def save(self, path: Union[str, os.PathLike]) -> int:
         """Saves the set of exported programs to a descriptor file.
@@ -86,7 +93,9 @@ def permute_path(name):
         count_deduped = 0
 
         # Save each.
-        for program_name, ep in self.programs.items():
+        for program_name, export_def in self.programs.items():
+            ep = export_def.target
+            assert isinstance(ep, torch.export.ExportedProgram)
             # First validate the ep with normal rules, which we will then
             # disable since we are violating the spec.
             ep._validate()
@@ -129,7 +138,7 @@ def load(path: Union[str, os.PathLike]) -> "FxPrograms":
             ep = torch.export.load(path.parent / program_file_name)
             _unsharify_state_dict(shared_state_dict, ep.state_dict)
             _unsharify_state_dict(shared_constants, _get_optional_constants(ep))
-            instance.programs[program_name] = ep
+            instance.programs[program_name] = ExportTargetDef(ep)
         return instance
 
 
@@ -169,6 +178,7 @@ def export_program(
         dynamic_shapes=None,
         strict: bool = True,
         name: Optional[str] = None,
+        arg_device: dict[int, DeviceAffinity] | None = None,
     ):
         if f is None:
             return functools.partial(
@@ -178,6 +188,7 @@ def export_program(
                 strict=strict,
                 dynamic_shapes=dynamic_shapes,
                 name=name,
+                arg_device=arg_device,
             )
 
         if name is None:
@@ -234,7 +245,10 @@ def new_forward(self, *forward_args, **forward_kwargs):
 
             _patch_op_dispatch_for_export()
             program = program.run_decompositions(current_decomps)
-        fx_builder.programs[name] = program
+        fx_builder.programs[name] = ExportTargetDef(
+            program,
+            arg_device=arg_device,
+        )
         return program