Add device affinities for arguments in AOT #231

sogartar · 2024-10-18T16:01:07Z

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.

sogartar · 2024-10-18T19:41:24Z

iree/turbine/aot/compiled_module.py

@@ -107,12 +112,27 @@ def __call__(self, *args, **kwargs):
        return self.py_value(*args, **kwargs)


+class ExportTargetDef:


I am not sure if I should drop the ExportTargetDef and use half-initialized ExportProcDef and ExportedProgramDef directly.

What is the benefit to having this separate class?

In the place where we need a structure to store this data we are not ready to construct ExportProcDef or ExportedProgramDef. If we are to use them directly they will have a more complicated multi-step initialization.

sogartar · 2024-10-18T19:43:09Z

This PR removes the need for #220, which kind of abuses the function generation mechanism.

stellaraccident · 2024-10-18T20:21:23Z

Thanks. I will review this in a few minutes

rsuderman · 2024-10-22T17:51:58Z

iree/turbine/aot/compiled_module.py

+        argument_device_affinities: dict[int, "DeviceAffinity"] | None = None,
+    ):
+        self.target = target
+        self.argument_device_affinities = argument_device_affinities


Make the name more succinct. Instead of argument_device_affinities just arg_device. It should focus on where the argument is placed.

rsuderman · 2024-10-22T17:54:53Z

iree/turbine/aot/compiled_module.py

@@ -207,6 +240,21 @@ def globals_defs(self) -> Generator[Tuple[str, GlobalsDef], None, None]:
        )  # type: ignore

    def def_attribute(self, key, value):
+        if isinstance(value, ExportTargetDef):
+            if isinstance(value.target, ExportedProgram):


Swap the if else and remove the else: part. You can do

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.argument_device_affinities )

And given the return it would exit right away anyway.

rsuderman · 2024-10-22T17:56:35Z

iree/turbine/aot/compiled_module.py

@@ -633,6 +708,7 @@ def __new__(
                ep_def.exported_program,
                symbol_name=ep_def.export_name or "main",
                symbol_visibility=None if ep_def.public else "private",
+                argument_device_affinities=ep_def.argument_device_affinities or {},


We should just support whatever the default is for ep_def.argument_device_affinities rather than using or

rsuderman · 2024-10-22T17:58:50Z

tests/aot/fx_programs_test.py

Do not wipe out the existing test. Create a separate fx_programs_test_device.py that tests the device affinity work. We should try to guarantee the old patch works for as long as possible

I moved the new test to a new file as you suggested.

rsuderman · 2024-10-22T18:12:57Z

iree/turbine/aot/compiled_module.py

@@ -207,6 +240,21 @@ def globals_defs(self) -> Generator[Tuple[str, GlobalsDef], None, None]:
        )  # type: ignore

    def def_attribute(self, key, value):
+        if isinstance(value, ExportTargetDef):
+            if isinstance(value.target, ExportedProgram):
+                value = ExportedProgramDef(


Shouldn't this be returned as well? It is setting the value but this falls through. I see it is handled further down, if so we should handle it right before the follow up case.

I moved this case down before the handling of ExportedProgramDef.

rsuderman · 2024-10-22T18:14:09Z

iree/turbine/aot/exporter.py

+class DeviceAffinity:
+    """This is used to provide device affinities to exported function arguments."""
+
+    def __init__(self, moniker: str):


Make the moniker an int. All cases around just specify it that way anyway.

I changed it to ordinal: int.

rsuderman · 2024-10-22T18:19:25Z

iree/turbine/aot/compiled_module.py

@@ -568,6 +627,22 @@ def save_mlir(inst: "CompiledModule", path: Union[Path, str]):

    jittable = staticmethod(builtins.jittable)

+    @staticmethod
+    def annotate(


Name this something better than annotate. Even signature_info would be sufficient to explain that its adding additional information to the function signature.

sogartar · 2024-10-22T19:13:57Z

iree/turbine/aot/fx_programs.py

@@ -61,7 +68,7 @@ class FxPrograms:
    """

    def __init__(self):
-        self.programs: dict[str, torch.export.ExportedProgram] = {}
+        self.programs: dict[str, ExportTargetDef] = {}


@rsuderman This is something I forgot to bring to attention, but I am not sure if self.programs should be a part of the interface. This changes it and I also changed one test that specifically used it.

stellaraccident

Thanks. A minor comment on code organization.

stellaraccident · 2024-10-22T21:29:17Z

iree/turbine/aot/exporter.py

@@ -49,6 +50,21 @@
 SaveableTarget = Union[str, Path, None, Output]


+class DeviceAffinity:


Put this in tensor_traits.py. Then you won't need to work around weird circular references and can just cleanly import it from anywhere needed.

I moved it there. Why I did not see this.

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 <[email protected]>

sogartar · 2024-10-23T13:19:58Z

I squashed and rebased to prepare for merging.

sogartar force-pushed the affinity-attributes-in-api branch 7 times, most recently from ce0da33 to 00d8857 Compare October 18, 2024 19:24

sogartar requested review from stellaraccident and rsuderman October 18, 2024 19:36

sogartar changed the title ~~WIP Add device affinities for arguments in AOT~~ Add device affinities for arguments in AOT Oct 18, 2024

sogartar marked this pull request as ready for review October 18, 2024 19:36

sogartar commented Oct 18, 2024

View reviewed changes

sogartar mentioned this pull request Oct 21, 2024

Add exporting with argument device affinities nod-ai/shark-ai#300

Merged

rsuderman requested changes Oct 22, 2024

View reviewed changes

sogartar commented Oct 22, 2024

View reviewed changes

sogartar requested a review from rsuderman October 22, 2024 19:14

stellaraccident approved these changes Oct 22, 2024

View reviewed changes

sogartar force-pushed the affinity-attributes-in-api branch from 9baff6b to a7e4bc4 Compare October 23, 2024 13:19

rsuderman approved these changes Oct 23, 2024

View reviewed changes

sogartar merged commit ae9a51c into iree-org:main Oct 23, 2024
7 of 8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add device affinities for arguments in AOT #231

Add device affinities for arguments in AOT #231

sogartar commented Oct 18, 2024

sogartar Oct 18, 2024

dan-garvey Oct 21, 2024

sogartar Oct 21, 2024

sogartar commented Oct 18, 2024

stellaraccident commented Oct 18, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

rsuderman Oct 22, 2024

sogartar Oct 22, 2024

sogartar Oct 22, 2024

stellaraccident left a comment

stellaraccident Oct 22, 2024

sogartar Oct 23, 2024

sogartar commented Oct 23, 2024

		@@ -107,12 +112,27 @@ def __call__(self, args, *kwargs):
		return self.py_value(args, *kwargs)


		class ExportTargetDef:

		@@ -49,6 +50,21 @@
		SaveableTarget = Union[str, Path, None, Output]


		class DeviceAffinity:

Add device affinities for arguments in AOT #231

Add device affinities for arguments in AOT #231

Conversation

sogartar commented Oct 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sogartar commented Oct 18, 2024

stellaraccident commented Oct 18, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stellaraccident left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sogartar commented Oct 23, 2024