| #!/usr/bin/env python3 | |
| # Copyright 2022 The IREE Authors | |
| # | |
| # Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| # See https://llvm.org/LICENSE.txt for license information. | |
| # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| import unittest | |
| from e2e_test_framework.definitions import common_definitions, iree_definitions | |
| import export_benchmark_config | |
| COMMON_MODEL = common_definitions.Model( | |
| id="tflite", | |
| name="model_tflite", | |
| tags=[], | |
| source_type=common_definitions.ModelSourceType.EXPORTED_TFLITE, | |
| source_url="", | |
| entry_function="predict", | |
| input_types=["1xf32"]) | |
| COMMON_GEN_CONFIG = iree_definitions.ModuleGenerationConfig( | |
| imported_model=iree_definitions.ImportedModel.from_model(COMMON_MODEL), | |
| compile_config=iree_definitions.CompileConfig(id="1", | |
| tags=[], | |
| compile_targets=[])) | |
| COMMON_EXEC_CONFIG = iree_definitions.ModuleExecutionConfig( | |
| id="exec", | |
| tags=[], | |
| loader=iree_definitions.RuntimeLoader.EMBEDDED_ELF, | |
| driver=iree_definitions.RuntimeDriver.LOCAL_SYNC) | |
| class ExportBenchmarkConfigTest(unittest.TestCase): | |
| def test_filter_and_group_run_configs_set_all_filters(self): | |
| device_spec_a = common_definitions.DeviceSpec( | |
| id="dev_a_cpu", | |
| device_name="dev_a_cpu", | |
| architecture=common_definitions.DeviceArchitecture.RV64_GENERIC, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| device_spec_b = common_definitions.DeviceSpec( | |
| id="dev_a_gpu", | |
| device_name="dev_a_gpu", | |
| architecture=common_definitions.DeviceArchitecture.MALI_VALHALL, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| device_spec_c = common_definitions.DeviceSpec( | |
| id="dev_c", | |
| device_name="dev_c", | |
| architecture=common_definitions.DeviceArchitecture.CUDA_SM80, | |
| host_environment=common_definitions.HostEnvironment.LINUX_X86_64) | |
| matched_run_config_a = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_a, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| unmatched_run_config_b = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_b, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| matched_run_config_c = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_c, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| matchers = [(lambda config: config.target_device_spec.architecture. | |
| architecture == "cuda"), | |
| (lambda config: config.target_device_spec.host_environment. | |
| platform == "android")] | |
| run_config_map = export_benchmark_config.filter_and_group_run_configs( | |
| run_configs=[ | |
| matched_run_config_a, unmatched_run_config_b, matched_run_config_c | |
| ], | |
| target_device_names={"dev_a_cpu", "dev_c"}, | |
| preset_matchers=matchers) | |
| self.assertEqual(run_config_map, { | |
| "dev_a_cpu": [matched_run_config_a], | |
| "dev_c": [matched_run_config_c], | |
| }) | |
| def test_filter_and_group_run_configs_include_all(self): | |
| device_spec_a = common_definitions.DeviceSpec( | |
| id="dev_a_cpu", | |
| device_name="dev_a_cpu", | |
| architecture=common_definitions.DeviceArchitecture.RV64_GENERIC, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| device_spec_b = common_definitions.DeviceSpec( | |
| id="dev_a_gpu", | |
| device_name="dev_a_gpu", | |
| architecture=common_definitions.DeviceArchitecture.MALI_VALHALL, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| device_spec_c = common_definitions.DeviceSpec( | |
| id="dev_a_second_gpu", | |
| device_name="dev_a_gpu", | |
| architecture=common_definitions.DeviceArchitecture.ADRENO_GENERIC, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| run_config_a = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_a, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_b = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_b, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_c = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_c, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_map = export_benchmark_config.filter_and_group_run_configs( | |
| run_configs=[run_config_a, run_config_b, run_config_c]) | |
| self.maxDiff = 100000 | |
| self.assertEqual(run_config_map, { | |
| "dev_a_cpu": [run_config_a], | |
| "dev_a_gpu": [run_config_b, run_config_c], | |
| }) | |
| def test_filter_and_group_run_configs_set_target_device_names(self): | |
| device_spec_a = common_definitions.DeviceSpec( | |
| id="dev_a", | |
| device_name="dev_a", | |
| architecture=common_definitions.DeviceArchitecture.RV64_GENERIC, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| device_spec_b = common_definitions.DeviceSpec( | |
| id="dev_b", | |
| device_name="dev_b", | |
| architecture=common_definitions.DeviceArchitecture.MALI_VALHALL, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| run_config_a = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_a, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_b = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=COMMON_GEN_CONFIG, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_b, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_map = export_benchmark_config.filter_and_group_run_configs( | |
| run_configs=[run_config_a, run_config_b], | |
| target_device_names={"dev_a", "dev_b"}) | |
| self.assertEqual(run_config_map, { | |
| "dev_a": [run_config_a], | |
| "dev_b": [run_config_b], | |
| }) | |
| def test_filter_and_group_run_configs_set_preset_matchers(self): | |
| small_model = common_definitions.Model( | |
| id="small_model", | |
| name="small_model", | |
| tags=[], | |
| source_type=common_definitions.ModelSourceType.EXPORTED_TFLITE, | |
| source_url="", | |
| entry_function="predict", | |
| input_types=["1xf32"]) | |
| big_model = common_definitions.Model( | |
| id="big_model", | |
| name="big_model", | |
| tags=[], | |
| source_type=common_definitions.ModelSourceType.EXPORTED_TFLITE, | |
| source_url="", | |
| entry_function="predict", | |
| input_types=["1xf32"]) | |
| compile_config = iree_definitions.CompileConfig(id="1", | |
| tags=[], | |
| compile_targets=[]) | |
| small_gen_config = iree_definitions.ModuleGenerationConfig( | |
| imported_model=iree_definitions.ImportedModel.from_model(small_model), | |
| compile_config=compile_config) | |
| big_gen_config = iree_definitions.ModuleGenerationConfig( | |
| imported_model=iree_definitions.ImportedModel.from_model(big_model), | |
| compile_config=compile_config) | |
| device_spec_a = common_definitions.DeviceSpec( | |
| id="dev_a", | |
| device_name="dev_a", | |
| architecture=common_definitions.DeviceArchitecture.RV64_GENERIC, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| device_spec_b = common_definitions.DeviceSpec( | |
| id="dev_b", | |
| device_name="dev_b", | |
| architecture=common_definitions.DeviceArchitecture.MALI_VALHALL, | |
| host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A) | |
| run_config_a = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=small_gen_config, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_a, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_b = iree_definitions.E2EModelRunConfig( | |
| module_generation_config=big_gen_config, | |
| module_execution_config=COMMON_EXEC_CONFIG, | |
| target_device_spec=device_spec_b, | |
| input_data=common_definitions.ZEROS_MODEL_INPUT_DATA) | |
| run_config_map = export_benchmark_config.filter_and_group_run_configs( | |
| run_configs=[run_config_a, run_config_b], | |
| preset_matchers=[ | |
| lambda config: config.module_generation_config.imported_model.model. | |
| id == "small_model" | |
| ]) | |
| self.assertEqual(run_config_map, { | |
| "dev_a": [run_config_a], | |
| }) | |
| if __name__ == "__main__": | |
| unittest.main() |