Add Flux_1.dev transformer and text encoder tests (nod-ai#411)

The most significant infrastructure change is with some reworks to
OnnxModelInfo:

1. `OnnxModelInfo` now has an attribute for `ExtraOptions`, which holds
test-specific options for later testing stages beyond setup. This was
required to allow running with external params, passing compile flags,
and setting custom mlir import options.
2. All attributes configurable via `update_*` methods in `OnnxModelInfo`
are updated in the init method immediately. This required reworking a
few helper classes which relied on calling these at a specific time.

Another infrastructure change is through the migration to using
`iree-import-onnx` rather than torch-mlir's onnx importer. This is to
allow externalizing parameters.

Adds the following three tests:

- `flux_1_dev_transformer`
- `flux_1_dev_clip`
- `flux_1_dev_t5`

Author: zjgarvey

alt_e2eshark/base_requirements.txt

@@ -6,6 +6,7 @@ ml_dtypes
 onnx
 onnxruntime
 transformers
+diffusers
 huggingface-hub
 sentencepiece
 accelerate

alt_e2eshark/e2e_testing/backends.py

@@ -7,7 +7,7 @@
 import onnxruntime as ort
 from typing import TypeVar, List
 from e2e_testing.storage import TestTensors, get_shape_string
-from e2e_testing.framework import CompiledOutput, ModelArtifact
+from e2e_testing.framework import CompiledOutput, ModelArtifact, CompilerOptions, RuntimeOptions
 from onnx import ModelProto
 import os
 from pathlib import Path
@@ -18,45 +18,56 @@
 class BackendBase(abc.ABC):
 
     @abc.abstractmethod
-    def compile(self, module: ModelArtifact) -> CompiledOutput:
+    def compile(self, module: ModelArtifact, extra_options : CompilerOptions) -> CompiledOutput:
         """specifies how to compile an MLIR Module"""
 
     @abc.abstractmethod
-    def load(self, artifact: CompiledOutput, func_name: str) -> Invoker:
+    def load(self, artifact: CompiledOutput, func_name: str, extra_options : RuntimeOptions) -> Invoker:
         """loads the function with name func_name from compiled artifact. This method should return a function callable from python."""
 
 
 from iree import compiler as ireec
 from iree import runtime as ireert
 
 
+def flag(arg : str) -> str:
+    if arg.startswith("--"):
+        return arg
+    return f'--{arg}'
+
 class SimpleIREEBackend(BackendBase):
     '''This backend uses iree to compile and run MLIR modules for a specified hal_target_backend'''
     def __init__(self, *, device="local-task", hal_target_backend="llvm-cpu", extra_args : List[str] = None):
         self.device = device
         self.hal_target_backend = hal_target_backend
-        self.extra_args = []
-        if extra_args:
-            for a in extra_args:
-                if a[0:2] == "--":
-                    self.extra_args.append(a)
-                else:
-                    self.extra_args.append("--" + a)
-
-    def compile(self, module, *, save_to: str = None):
+        self.extra_args = [] if extra_args is None else [flag(a) for a in extra_args]
+        if hal_target_backend == "rocm":
+            self.extra_args += [
+                f"--iree-hip-target={self.target_chip}",
+            ]
+        if hal_target_backend == "llvm-cpu":
+            self.extra_args += [
+                "--iree-llvmcpu-target-cpu=host",
+            ]
+
+    def compile(self, module, *, save_to: str = None, extra_options : CompilerOptions):
+        test_specific_args = list(extra_options.common_extra_args)
+        if self.hal_target_backend in extra_options.backend_specific_flags.keys():
+            test_specific_args += list(extra_options.backend_specific_flags[self.hal_target_backend])
+        compile_args = self.extra_args + [flag(arg) for arg in test_specific_args]
         # compile to a vmfb for llvm-cpu
         b = ireec.tools.compile_str(
             str(module),
             target_backends=[self.hal_target_backend],
-            extra_args=self.extra_args,
+            extra_args=compile_args,
         )
         # log the vmfb
         if save_to:
             with open(os.path.join(save_to, "compiled_model.vmfb"), "wb") as f:
                 f.write(b)
         return b
 
-    def load(self, artifact, *, func_name="main"):
+    def load(self, artifact, *, func_name="main", extra_options : RuntimeOptions):
         config = ireert.Config(self.device)
         ctx = ireert.SystemContext(config=config)
         vm_module = ireert.VmModule.copy_buffer(ctx.instance, artifact)
@@ -80,13 +91,7 @@ def __init__(self, *, device="local-task", hal_target_backend="llvm-cpu", target
         self.device = device
         self.hal_target_backend = hal_target_backend
         self.target_chip = target_chip
-        self.extra_args = []
-        if extra_args:
-            for a in extra_args:
-                if a[0:2] == "--":
-                    self.extra_args.append(a)
-                else:
-                    self.extra_args.append("--" + a)
+        self.extra_args = [] if extra_args is None else [flag(a) for a in extra_args]
         if hal_target_backend == "rocm":
             self.extra_args += [
                 f"--iree-hip-target={self.target_chip}",
@@ -96,15 +101,17 @@ def __init__(self, *, device="local-task", hal_target_backend="llvm-cpu", target
                 "--iree-llvmcpu-target-cpu=host",
             ]
 
-    def compile(self, module_path: str, *, save_to : str = None) -> str:
+    def compile(self, module_path: str, *, save_to : str = None, extra_options : CompilerOptions) -> str:
+        test_specific_args = list(extra_options.common_extra_args)
+        if self.hal_target_backend in extra_options.backend_specific_flags.keys():
+            test_specific_args += list(extra_options.backend_specific_flags[self.hal_target_backend])
+        compile_args = self.extra_args + [flag(arg) for arg in test_specific_args]
         vmfb_path = os.path.join(save_to, "compiled_model.vmfb")
         arg_string = f"--iree-hal-target-backends={self.hal_target_backend} "
-        for arg in self.extra_args:
-            arg_string += arg
-            arg_string += " "
+        arg_string += ' '.join(compile_args)
         detail_log = os.path.join(save_to, "detail", "compilation.detail.log")
         commands_log = os.path.join(save_to, "commands", "compilation.commands.log")
-        script = f"iree-compile {module_path} {arg_string}-o {vmfb_path} 1> {detail_log} 2>&1"
+        script = f"iree-compile {module_path} {arg_string} -o {vmfb_path} 1> {detail_log} 2>&1"
         with open(commands_log, "w") as file:
             file.write(script) 
         # remove old vmfb if it exists
@@ -116,16 +123,21 @@ def compile(self, module_path: str, *, save_to : str = None) -> str:
             raise FileNotFoundError(error_msg)
         return vmfb_path
 
-    def load(self, vmfb_path: str, *, func_name=None):
+    def load(self, vmfb_path: str, *, func_name=None, extra_options : RuntimeOptions):
         """A bit hacky. func returns a script that would dump outputs to terminal output. Modified in config.run method"""
+        test_specific_args = list(extra_options.common_extra_args)
+        if self.hal_target_backend in extra_options.backend_specific_flags.keys():
+            test_specific_args += list(extra_options.backend_specific_flags[self.hal_target_backend])
         run_dir = Path(vmfb_path).parent
         def func(x: TestTensors) -> str:
-            script = f"iree-run-module --module='{vmfb_path}' --device={self.device}"
+            script = f"iree-run-module --module='{vmfb_path}' --device={self.device} "
+            for arg in test_specific_args:
+                script += f'{flag(arg)} '
             if func_name:
-                script += f" --function='{func_name}'"
+                script += f"--function='{func_name}' "
             torch_inputs = x.to_torch().data
             for index, input in enumerate(torch_inputs):
-                script += f" --input='{get_shape_string(input)}=@{run_dir}/input.{index}.bin'"
+                script += f"--input='{get_shape_string(input)}=@{run_dir}/input.{index}.bin' "
             return script
         return func
 
@@ -135,16 +147,10 @@ class OnnxrtIreeEpBackend(BackendBase):
     def __init__(self, *, device="local-task", hal_target_device="llvm-cpu", extra_args : List[str] = None):
         self.device = device
         self.hal_target_device = hal_target_device
-        if extra_args:
-            self.extra_args = []
-            for a in extra_args:
-                if a[0:2] == "--":
-                    self.extra_args.append(a)
-                else:
-                    self.extra_args.append("--" + a)
-        elif hal_target_device == "hip":
+        self.extra_args = [] if extra_args is None else [flag(a) for a in extra_args]
+        if hal_target_device == "hip":
             # some extra args for Mi250 - some of these may not work for other chips
-            self.extra_args = [
+            self.extra_args += [
                 "--iree-hip-target=gfx90a",
             ]
         self.providers = ["IreeExecutionProvider"]
@@ -159,7 +165,7 @@ def __init__(self, *, device="local-task", hal_target_device="llvm-cpu", extra_a
         #  sess_opt.log_verbosity_level = 0
         #  self.sess_opt.log_severity_level = 0
 
-    def compile(self, model: ModelProto, *, save_to: str = None) -> ort.InferenceSession:
+    def compile(self, model: ModelProto, *, save_to: str = None, extra_options : CompilerOptions) -> ort.InferenceSession:
         if self.provider_options:
             provider_options_dict = self.provider_options[0]
             provider_options_dict["save_to"] = save_to
@@ -173,7 +179,7 @@ def compile(self, model: ModelProto, *, save_to: str = None) -> ort.InferenceSes
         # can't save an onnx runtime session
         return session
 
-    def load(self, session: ort.InferenceSession, *, func_name=None) -> Invoker:
+    def load(self, session: ort.InferenceSession, *, func_name=None, extra_options : RuntimeOptions) -> Invoker:
         def func(x: TestTensors):
             data = x.to_numpy().data
             session_inputs = session.get_inputs()

alt_e2eshark/e2e_testing/framework.py

@@ -8,14 +8,39 @@
 import abc
 import os
 from pathlib import Path
-from typing import Union, TypeVar, Tuple, NamedTuple, Dict, Optional, Callable
+from typing import Union, TypeVar, Tuple, NamedTuple, Dict, Optional, Callable, List
 from e2e_testing.storage import TestTensors
 from e2e_testing.onnx_utils import *
 
 # This file two types of classes: framework-specific base classes for storing model info, and generic classes for testing infrastructure.
 
 Module = TypeVar("Module")
 
+class ImporterOptions(NamedTuple):
+    opset_version : Optional[int] = None
+    large_model : bool = False
+    externalize_params : bool = False
+    externalize_inputs_threshold : Optional[int] = None
+    num_elements_threshold: int = 100
+    params_scope : str = "model"
+    param_gb_threshold : Optional[float] = None
+
+class CompilerOptions(NamedTuple):
+    """Specify, for specific iree-hal-target-backends, a tuple of extra compiler flags. 
+       Also allows backend-agnostic options to be included."""
+    backend_specific_flags : Dict[str, Tuple[str]] = dict()
+    common_extra_args : Tuple[str] = tuple()
+
+class RuntimeOptions(NamedTuple):
+    """Specify, for specific iree-hal-target-backends, a tuple of extra runtime flags.
+       Also allows backend-agnostic options to be included."""
+    backend_specific_flags : Dict[str, Tuple[str]] = dict()
+    common_extra_args : Tuple[str] = tuple()
+
+class ExtraOptions(NamedTuple):
+    import_model_options : ImporterOptions = ImporterOptions()
+    compilation_options : CompilerOptions = CompilerOptions()
+    compiled_inference_options : RuntimeOptions = RuntimeOptions()
 
 class OnnxModelInfo:
     """Stores information about an onnx test: the filepath to model.onnx, how to construct/download it, and how to construct sample inputs for a test run."""
@@ -29,16 +54,21 @@ def __init__(
         self.name = name
         self.model = os.path.join(onnx_model_path, "model.onnx")
         self.opset_version = opset_version
-        self.sess_options = ort.SessionOptions()
+
         self.dim_param_dict = None
+        self.update_dim_param_dict()
         self.input_name_to_shape_map = None
+        self.update_input_name_to_shape_map()
+        self.sess_options = ort.SessionOptions()
+        self.update_sess_options()
+        self.extra_options = ExtraOptions()
+        self.update_extra_options()
 
     def forward(self, input: Optional[TestTensors] = None) -> TestTensors:
         """Applies self.model to self.input. Only override if necessary for specific models"""
         input = input.to_numpy().data
         if not os.path.exists(self.model):
             self.construct_model()
-        self.update_sess_options()
         session = ort.InferenceSession(self.model, self.sess_options)
         session_inputs = session.get_inputs()
         session_outputs = session.get_outputs()
@@ -50,23 +80,27 @@ def forward(self, input: Optional[TestTensors] = None) -> TestTensors:
 
         return TestTensors(model_output)
 
-    def update_sess_options(self):
-        """Can be overridden to modify session options (self.sess_options) for gold inference.
-        It is sometimes useful to disable all optimizations, which can be done with:
-        self.sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
-        """
-        pass
-
     def update_dim_param_dict(self):
         """Can be overridden to modify a dictionary of dim parameters (self.dim_param_dict) used to 
         construct inputs for a model with dynamic dims.
         """
         pass
 
-    def contruct_input_name_to_shape_map(self):
+    def update_input_name_to_shape_map(self):
         """Can be overriden to construct an assocation map between the name of the input nodes and their shapes."""
         pass
 
+    def update_sess_options(self):
+        """Can be overridden to modify session options (self.sess_options) for gold inference.
+        It is sometimes useful to disable all optimizations, which can be done with:
+        self.sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+        """
+        pass
+
+    def update_extra_options(self):
+        """Can be overridden to set self.extra_options = ExtraOptions(**kwargs)"""
+        pass
+
     def construct_model(self):
         """a method to be overwritten. To make a new test, define a subclass with an override for this method"""
         raise NotImplementedError(
@@ -151,7 +185,7 @@ def get_metadata(self):
 class TestConfig(abc.ABC):
 
     @abc.abstractmethod
-    def import_model(self, program: TestModel, *, save_to: str) -> Tuple[ModelArtifact, str | None]:
+    def import_model(self, program: TestModel, *, save_to: str, extra_options : ImporterOptions) -> Tuple[ModelArtifact, str | None]:
         """imports the test model to model artifact (e.g., loads the onnx model )"""
         pass
 
@@ -161,16 +195,16 @@ def preprocess_model(self, model_artifact: ModelArtifact, *, save_to: str) -> Mo
         pass
 
     @abc.abstractmethod
-    def compile(self, module: ModelArtifact, *, save_to: str) -> CompiledOutput:
+    def compile(self, module: ModelArtifact, *, save_to: str, extra_options : CompilerOptions) -> CompiledOutput:
         """converts the test program to a compiled artifact"""
         pass
 
     @abc.abstractmethod
-    def run(self, artifact: CompiledOutput, input: TestTensors) -> TestTensors:
+    def run(self, artifact: CompiledOutput, input: TestTensors, extra_options : RuntimeOptions) -> TestTensors:
         """runs the input through the compiled artifact"""
         pass
 
-    def benchmark(self, artifact: CompiledOutput, input: TestTensors, repetitions: int, *, func_name=None) -> float:
+    def benchmark(self, artifact: CompiledOutput, input: TestTensors, repetitions: int, *, func_name=None, extra_options : RuntimeOptions) -> float:
         """returns a float representing inference time in ms"""
         pass
 

alt_e2eshark/e2e_testing/storage.py

@@ -171,20 +171,14 @@ def to_dtype(self, dtype, *, index: Optional[int] = None) -> "TestTensors":
         """returns a copy of self with a converted dtype (at a particular index, if specified)"""
         if self.type == numpy.ndarray:
             if index:
-                try:
-                    new_data = self.data
-                    new_data[index] = new_data[index].astype(dtype)
-                except Exception as e:
-                    print("to_dtype failed due to excepton {e}.")
+                new_data = self.data
+                new_data[index] = new_data[index].astype(dtype)
             else:
                 new_data = tuple([d.astype(dtype) for d in self.data])
         if self.type == torch.Tensor:
             if index:
-                try:
-                    new_data = self.data
-                    new_data[index] = new_data[index].to(dtype=dtype)
-                except Exception as e:
-                    print("to_dtype failed due to excepton {e}.")
+                new_data = self.data
+                new_data[index] = new_data[index].to(dtype=dtype)
             else:
                 new_data = tuple([d.to(dtype=dtype) for d in self.data])
         return TestTensors(new_data)