Add TF GPT2 benchmarks

nod-ai · Oct 12, 2023 · aef639b · aef639b
1 parent 4e06f74
commit aef639b
Show file tree

Hide file tree

Showing 8 changed files with 667 additions and 14 deletions.
diff --git a/build_tools/python/benchmark_suites/iree/adreno_benchmarks.py b/build_tools/python/benchmark_suites/iree/adreno_benchmarks.py
@@ -10,7 +10,7 @@
 from benchmark_suites.iree import benchmark_presets, module_execution_configs, utils
 from e2e_test_framework import unique_ids
 from e2e_test_framework.definitions import common_definitions, iree_definitions
-from e2e_test_framework.models import tflite_models
+from e2e_test_framework.models import tflite_models, tf_models
 from e2e_test_framework.device_specs import device_collections
 
 
@@ -46,6 +46,8 @@ def generate(
     ) -> List[iree_definitions.E2EModelRunConfig]:
         default_models = [
             tflite_models.MOBILEBERT_FP32,
+            tf_models.GPT2_117M_1x4_FP32_TF,
+            tf_models.GPT2_117M_1x1_FP32_TF,
         ]
         default_gen_configs = [
             iree_definitions.ModuleGenerationConfig.build(

diff --git a/build_tools/python/benchmark_suites/iree/armv8_a_benchmarks.py b/build_tools/python/benchmark_suites/iree/armv8_a_benchmarks.py
@@ -11,7 +11,7 @@
 from e2e_test_framework import unique_ids
 from e2e_test_framework.definitions import common_definitions, iree_definitions
 from e2e_test_framework.device_specs import device_collections
-from e2e_test_framework.models import tflite_models
+from e2e_test_framework.models import tflite_models, tf_models
 
 
 class Android_ARMv8_A_Benchmarks(object):
@@ -20,6 +20,8 @@ class Android_ARMv8_A_Benchmarks(object):
     NONQUANT_MODELS = [
         tflite_models.DEEPLABV3_FP32,
         tflite_models.MOBILEBERT_FP32,
+        tf_models.GPT2_117M_1x4_FP32_TF,
+        tf_models.GPT2_117M_1x1_FP32_TF,
     ]
     QUANT_MODELS = [tflite_models.MOBILEBERT_INT8]
 
@@ -88,10 +90,6 @@ def generate(
             for model in self.QUANT_MODELS
         ]
 
-        all_devices = device_collections.DEFAULT_DEVICE_COLLECTION.query_device_specs(
-            architecture=common_definitions.DeviceArchitecture.ARMV8_2_A_GENERIC,
-            host_environment=common_definitions.HostEnvironment.ANDROID_ARMV8_2_A,
-        )
         big_cores_devices = (
             device_collections.DEFAULT_DEVICE_COLLECTION.query_device_specs(
                 architecture=common_definitions.DeviceArchitecture.ARMV8_2_A_GENERIC,
@@ -103,13 +101,7 @@ def generate(
             module_generation_configs=default_gen_confings,
             module_execution_configs=local_sync_execution_configs
             + local_task_execution_configs,
-            device_specs=all_devices,
-            presets=[benchmark_presets.ANDROID_CPU],
-        )
-        run_configs += utils.generate_e2e_model_run_configs(
-            module_generation_configs=experimental_gen_confings,
-            module_execution_configs=local_sync_execution_configs,
-            device_specs=all_devices,
+            device_specs=big_cores_devices,
             presets=[benchmark_presets.ANDROID_CPU],
         )
         run_configs += utils.generate_e2e_model_run_configs(

diff --git a/build_tools/python/benchmark_suites/iree/mali_benchmarks.py b/build_tools/python/benchmark_suites/iree/mali_benchmarks.py
@@ -10,7 +10,7 @@
 from benchmark_suites.iree import benchmark_presets, module_execution_configs, utils
 from e2e_test_framework import unique_ids
 from e2e_test_framework.definitions import common_definitions, iree_definitions
-from e2e_test_framework.models import tflite_models
+from e2e_test_framework.models import tflite_models, tf_models
 from e2e_test_framework.device_specs import device_collections
 
 
@@ -60,6 +60,8 @@ class Android_Mali_Benchmarks(object):
 
     FP32_MODELS = [
         tflite_models.MOBILEBERT_FP32,
+        tf_models.GPT2_117M_1x4_FP32_TF,
+        tf_models.GPT2_117M_1x1_FP32_TF,
     ]
     FP16_MODELS = [tflite_models.MOBILEBERT_FP16]
     QUANT_MODELS = [

diff --git a/build_tools/python/e2e_test_framework/models/model_groups.py b/build_tools/python/e2e_test_framework/models/model_groups.py
@@ -73,6 +73,12 @@
     common_definitions.CpuBenchmarkConfig(
         model=tf_models.BERT_LARGE_TF_FP32_SEQLEN384, threads=[8]
     ),
+    common_definitions.CpuBenchmarkConfig(
+        model=tf_models.GPT2_117M_1x4_FP32_TF, threads=[8]
+    ),
+    common_definitions.CpuBenchmarkConfig(
+        model=tf_models.GPT2_117M_1x1_FP32_TF, threads=[8]
+    ),
 ]
 
 # A subset of `x86_64_MODELS_AND_THREADS`.
@@ -120,6 +126,12 @@
     common_definitions.CpuBenchmarkConfig(
         model=tf_models.BERT_LARGE_TF_FP32_SEQLEN384, threads=[8]
     ),
+    common_definitions.CpuBenchmarkConfig(
+        model=tf_models.GPT2_117M_1x4_FP32_TF, threads=[8]
+    ),
+    common_definitions.CpuBenchmarkConfig(
+        model=tf_models.GPT2_117M_1x1_FP32_TF, threads=[8]
+    ),
 ]
 
 X86_64_BENCHMARK_CONFIG_LONG = [

diff --git a/build_tools/python/e2e_test_framework/models/tf_models.py b/build_tools/python/e2e_test_framework/models/tf_models.py
@@ -61,6 +61,26 @@
     input_types=["1x384xi32", "1x384xi32", "1x384xi32"],
 )
 
+GPT2_117M_1x4_FP32_TF = common_definitions.Model(
+    id=unique_ids.MODEL_GPT2_117M_1x4_FP32_TF,
+    name="GPT2_117M_TF_1X4XI32",
+    tags=["fp32", "tensorflow", "gpt2", "batch-1"],
+    source_type=common_definitions.ModelSourceType.EXPORTED_STABLEHLO_MLIR,
+    source_url="https://storage.googleapis.com/iree-shared-files/tf_gpt2/static_input_seqlen5/stablehlo.mlir",
+    entry_function="forward",
+    input_types=["1x4xi32", "1x4xi32"],
+)
+
+GPT2_117M_1x1_FP32_TF = common_definitions.Model(
+    id=unique_ids.MODEL_GPT2_117M_1x1_FP32_TF,
+    name="GPT2_117M_TF_1X1XI32",
+    tags=["fp32", "tensorflow", "gpt2", "batch-1"],
+    source_type=common_definitions.ModelSourceType.EXPORTED_STABLEHLO_MLIR,
+    source_url="https://storage.googleapis.com/iree-shared-files/tf_gpt2/static_input_seqlen1/stablehlo.mlir",
+    entry_function="forward",
+    input_types=["1x1xi32", "12x2x1x12x4x64xf32"],
+)
+
 TF_MODELS_ROOT_DIR = "https://storage.googleapis.com/iree-model-artifacts/tensorflow/tf_models_2.15.0.dev20230817_1692333975"
 
 ID_FORMAT = string.Template("${model_id}-batch-${batch_size}")

diff --git a/build_tools/python/e2e_test_framework/unique_ids.py b/build_tools/python/e2e_test_framework/unique_ids.py
@@ -85,6 +85,9 @@ def hash_composite_id(keys: Sequence[str]) -> str:
 MODEL_BERT_LARGE_384_FP32_TF = "5f3de3b3-fd00-4582-a97e-b70ff5edab07"
 MODEL_T5_LARGE_512_FP32_TF = "587e595d-2adf-4e41-9617-43178a133725"
 
+MODEL_GPT2_117M_1x4_FP32_TF = "52af9a47-e92c-4fbb-bdbf-6e1ee1a97e3a"
+MODEL_GPT2_117M_1x1_FP32_TF = "3f7398e6-bdb6-41ac-886e-0b3f17da603c"
+
 #    PyTorch.
 MODEL_CLIP_TEXT_SEQLEN64_FP32_TORCH = "9a9515c7-cb68-4c34-b1d2-0e8c0a3620b8"
 MODEL_UNET_2D_FP32_TORCH = "340553d1-e6fe-41b6-b2c7-687c74ccec56"

diff --git a/tests/e2e/test_artifacts/generated_e2e_test_fetch_models.cmake b/tests/e2e/test_artifacts/generated_e2e_test_fetch_models.cmake
@@ -110,6 +110,20 @@ iree_fetch_artifact(
   UNPACK
 )
 
+iree_fetch_artifact(
+  NAME "model-GPT2_117M_TF_1X4XI32"
+  SOURCE_URL "https://storage.googleapis.com/iree-shared-files/tf_gpt2/static_input_seqlen5/stablehlo.mlir"
+  OUTPUT "${ROOT_ARTIFACTS_DIR}/model_GPT2_117M_TF_1X4XI32.mlir"
+  UNPACK
+)
+
+iree_fetch_artifact(
+  NAME "model-GPT2_117M_TF_1X1XI32"
+  SOURCE_URL "https://storage.googleapis.com/iree-shared-files/tf_gpt2/static_input_seqlen1/stablehlo.mlir"
+  OUTPUT "${ROOT_ARTIFACTS_DIR}/model_GPT2_117M_TF_1X1XI32.mlir"
+  UNPACK
+)
+
 iree_fetch_artifact(
   NAME "model-BertLargeTFBatch1"
   SOURCE_URL "https://storage.googleapis.com/iree-model-artifacts/tensorflow/tf_models_2.15.0.dev20230817_1692333975/BERT_LARGE_FP32_TF_384XI32_BATCH1/stablehlo.mlirbc"