[WIP] Add InferenceBenchmark

iree-tf/benchmark/benchmark_model.py

@@ -15,7 +15,7 @@
 from typing import Optional
 
 # Add library dir to the search path.
-sys.path.insert(0, str(pathlib.Path(__file__).parent.parent / "library"))
+sys.path.insert(0, str(pathlib.Path(__file__).parents[1] / "library"))
 from models import resnet50, bert_large, t5_large
 
 # Add benchmark definitions to the search path.
@@ -46,6 +46,15 @@ def benchmark_lookup(unique_id: str):
     raise ValueError(f"Model definition not supported")
 
 
+def benchmark_lookup(benchmark_id: str):
+  benchmark = tf_inference_benchmarks.ID_TO_BENCHMARK_MAP.get(benchmark_id)
+  if benchmark is None:
+    raise ValueError(f"Id {benchmark_id} does not exist in benchmark suite.")
+
+  model_name, model_class = model_lookup(benchmark.model.id)
+  return (model_name, model_class, benchmark)
+
+
 def dump_result(file_path: str, result: dict) -> None:
   with open(file_path, "r") as f:
     dictObj = json.load(f)
@@ -66,7 +75,8 @@ def bytes_to_mb(bytes: Optional[int]) -> Optional[float]:
 def run_framework_benchmark(model_name: str, model_class: type[tf.Module],
                             batch_size: int, warmup_iterations: int,
                             benchmark_iterations: int, tf_device: str,
-                            hlo_dump_dir: str, dump_hlo: bool, shared_dict) -> None:
+                            hlo_dump_dir: str, dump_hlo: bool,
+                            shared_dict) -> None:
   try:
     with tf.device(tf_device):
       if dump_hlo:
@@ -216,17 +226,16 @@ def run_compiler_benchmark(hlo_benchmark_tool_path: str, hlo_dir: str,
 
   args = argParser.parse_args()
 
-  model_name, model_class, model_definition = benchmark_lookup(
-      args.benchmark_id)
+  model_name, model_class, benchmark = benchmark_lookup(args.benchmark_id)
   print(
-      f"\n\n--- {model_name} {args.benchmark_id} -------------------------------------"
+      f"\n\n--- {benchmark.name} {args.benchmark_id} -------------------------------------"
   )
 
   if os.path.exists(_HLO_DUMP_DIR):
     shutil.rmtree(_HLO_DUMP_DIR)
   os.mkdir(_HLO_DUMP_DIR)
 
-  batch_size = model_definition.input_batch_size
+  batch_size = benchmark.input_batch_size
   benchmark_definition = {
       "benchmark_id": args.benchmark_id,
       "benchmark_name": model_definition.name,
@@ -248,9 +257,9 @@ def run_compiler_benchmark(hlo_benchmark_tool_path: str, hlo_dir: str,
     shared_dict = manager.dict()
 
     if args.run_in_process:
-      run_framework_benchmark(model_name, model_class, batch_size, args.warmup_iterations,
-                              args.iterations, tf_device, _HLO_DUMP_DIR, dump_hlo,
-                              shared_dict)
+      run_framework_benchmark(model_name, model_class, batch_size,
+                              args.warmup_iterations, args.iterations,
+                              tf_device, _HLO_DUMP_DIR, dump_hlo, shared_dict)
     else:
       p = multiprocessing.Process(target=run_framework_benchmark,
                                   args=(model_name, model_class, batch_size,
@@ -269,8 +278,10 @@ def run_compiler_benchmark(hlo_benchmark_tool_path: str, hlo_dir: str,
       shared_dict = manager.dict()
 
       if args.run_in_process:
-        run_compiler_benchmark(args.hlo_benchmark_path, _HLO_DUMP_DIR, args.hlo_iterations,
-                              "cuda" if args.device == "gpu" else "cpu", shared_dict)
+        run_compiler_benchmark(args.hlo_benchmark_path, _HLO_DUMP_DIR,
+                               args.hlo_iterations,
+                               "cuda" if args.device == "gpu" else "cpu",
+                               shared_dict)
       else:
         p = multiprocessing.Process(
             target=run_compiler_benchmark,

oobi/benchmark-definitions/python/data_types.py

@@ -97,11 +97,32 @@ class Model(object):
   # Tags that describe the model characteristics.
   tags: List[str]
   meta_model: MetaModel
-  input_batch_size: int
-  inputs: ModelData
-  outputs: ModelData
   # A list of artifacts derived from this model.
   artifacts: List[ModelArtifact]
 
   def __str__(self):
     return self.name
+
+
+@serialization.serializable
+@dataclass(frozen=True)
+class InferenceBenchmark(object):
+  """Inference benchmark definition"""
+
+  id: str
+  # Unique friendly name.
+  name: str
+  # Tags that describe the benchmark characteristics.
+  tags: List[str]
+
+  # Model to inference
+  model: Model
+  # Model inputs
+  inputs: ModelData
+  # Model input batch size
+  input_batch_size: int
+  # Expected model outputs
+  outputs: ModelData
+
+  def __str__(self):
+    return self.name

oobi/benchmark-definitions/python/input_data_definitions.py

@@ -91,6 +91,15 @@
     ],
 )
 
+IMAGENET_APPLES_224X224X3XF32_BATCHES = {
+    1: IMAGENET_APPLES_224X224X3XF32_BATCH1,
+    8: IMAGENET_APPLES_224X224X3XF32_BATCH8,
+    64: IMAGENET_APPLES_224X224X3XF32_BATCH64,
+    128: IMAGENET_APPLES_224X224X3XF32_BATCH128,
+    256: IMAGENET_APPLES_224X224X3XF32_BATCH256,
+    2048: IMAGENET_APPLES_224X224X3XF32_BATCH2048,
+}
+
 BERT_LARGE_SEQLEN384_I32_BATCH1 = data_types.ModelData(
     id=unique_ids.INPUT_DATA_BERT_LARGE_SEQLEN384_I32_BATCH1,
     name="BERT_LARGE_SEQLEN384_I32_BATCH1",
@@ -244,6 +253,18 @@
     ],
 )
 
+BERT_LARGE_SEQLEN384_I32_BATCHES = {
+    1: BERT_LARGE_SEQLEN384_I32_BATCH1,
+    16: BERT_LARGE_SEQLEN384_I32_BATCH16,
+    24: BERT_LARGE_SEQLEN384_I32_BATCH24,
+    32: BERT_LARGE_SEQLEN384_I32_BATCH32,
+    48: BERT_LARGE_SEQLEN384_I32_BATCH48,
+    64: BERT_LARGE_SEQLEN384_I32_BATCH64,
+    512: BERT_LARGE_SEQLEN384_I32_BATCH512,
+    1024: BERT_LARGE_SEQLEN384_I32_BATCH1024,
+    1280: BERT_LARGE_SEQLEN384_I32_BATCH1280,
+}
+
 T5_LARGE_SEQLEN512_I32_BATCH1 = data_types.ModelData(
     id=unique_ids.INPUT_DATA_T5_LARGE_SEQLEN512_I32_BATCH1,
     name="T5_LARGE_SEQLEN512_I32_BATCH1",
@@ -369,3 +390,13 @@
         "https://storage.googleapis.com/iree-model-artifacts/tensorflow/tf_models_2.12.0_1681767794/T5_LARGE/batch_512/input_1.npy",
     ],
 )
+
+T5_LARGE_SEQLEN512_I32_BATCHES = {
+    1: T5_LARGE_SEQLEN512_I32_BATCH1,
+    16: T5_LARGE_SEQLEN512_I32_BATCH16,
+    24: T5_LARGE_SEQLEN512_I32_BATCH24,
+    32: T5_LARGE_SEQLEN512_I32_BATCH32,
+    48: T5_LARGE_SEQLEN512_I32_BATCH48,
+    64: T5_LARGE_SEQLEN512_I32_BATCH64,
+    512: T5_LARGE_SEQLEN512_I32_BATCH512,
+}