Move prepration of benchmark files to python script. (iree-org#5201)

This moves some hard-coded commands to a python script, so all the scripts can share informations in configuration.py. This makes flow easier when adding a new mako benchmark.

build_tools/buildkite/cmake/android/arm64-v8a/benchmark.yml

@@ -16,13 +16,8 @@ steps:
   - label: "build"
     commands:
       - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake-android@sha256:15d3266ae4865f7642a4ef4d76e5181f0dc3482a7cfba9021b6b55be524208ec build_tools/kokoro/gcp_ubuntu/cmake/android/build.sh arm64-v8a"
-      - "gsutil cp gs://iree-model-artifacts/iree-mobile-bert-artifacts-6fe4616e0ab9958eb18f368960a31276f1362029.tar.gz ."
-      - "gsutil cp gs://iree-model-artifacts/mobilenet-v2.tar.gz ."
-      - "tar -xzf iree-mobile-bert-artifacts-6fe4616e0ab9958eb18f368960a31276f1362029.tar.gz"
-      - "tar -xzf mobilenet-v2.tar.gz"
+      - "python3 build_tools/mako/prepare_benchmark_files.py"
       - "docker run --user=$(id -u):$(id -g) --volume=\\$PWD:\\$IREE_DOCKER_WORKDIR --workdir=\\$IREE_DOCKER_WORKDIR --rm gcr.io/iree-oss/cmake-android@sha256:15d3266ae4865f7642a4ef4d76e5181f0dc3482a7cfba9021b6b55be524208ec python3 build_tools/mako/compile_android_modules.py"
-      - "cp tmp/iree/modules/MobileBertSquad/iree_vmla/traces/serving_default/flagfile mobile-bert_flagfile"
-      - "cp mobilenet-v2/flagfile mobilenet-v2_flagfile"
       - "tar -czvf model-Pixel4-artifacts.tgz build-android/iree/tools/iree-benchmark-module *Pixel4*.vmfb"
       - "tar -czvf model-S20-artifacts.tgz build-android/iree/tools/iree-benchmark-module *S20*.vmfb"
       - "tar -czvf flagfiles.tgz *_flagfile"

build_tools/mako/benchmark_modules_on_android.py

@@ -88,7 +88,7 @@ def main(args) -> None:
       for target in phone.targets:
         module_name = configuration.get_module_name(model_benchmark.name,
                                                     phone.name, target.mako_tag)
-        flagfile_name = f"{model_benchmark.name}_flagfile"
+        flagfile_name = configuration.get_flagfile_name(model_benchmark.name)
         mako_log.append(benchmark(module_name, flagfile_name, target))
       mako_log.append(
           get_mako_metadata(args.git_hash, timestamp, phone.benchmark_key))

build_tools/mako/configuration.py

@@ -76,15 +76,21 @@ class ModelBenchmarkInfo:
 
   Attributes:
     name: The name of the model.
+    model_artifacts_name: The filename of model artifacts which locates in
+      Google bucket iree-model-artifacts.
     model_path: A path to MLIR input file. This can be a relative path.
+    flagfile_path: A path to flagfile. This can be a relative path.
     phones: A list of PhoneBenchmarkInfo that the benchmarking targets on.
   """
 
-  def __init__(self, name, model_path, phones):
+  def __init__(self, name, model_artifacts_name, model_path, flagfile_path,
+               phones):
     if "_" in name:
       raise ValueError("The target name contains invalid char '_'")
     self.name = name
+    self.model_artifacts_name = model_artifacts_name
     self.model_path = model_path
+    self.flagfile_path = flagfile_path
     self.phones = phones
 
 
@@ -144,7 +150,11 @@ def get_s20_default_target_list(batch_config=None):
 MODEL_BENCHMARKS = [
     ModelBenchmarkInfo(
         name="mobile-bert",
+        model_artifacts_name=
+        "iree-mobile-bert-artifacts-6fe4616e0ab9958eb18f368960a31276f1362029.tar.gz",
         model_path="tmp/iree/modules/MobileBertSquad/iree_input.mlir",
+        flagfile_path=
+        "tmp/iree/modules/MobileBertSquad/iree_vmla/traces/serving_default/flagfile",
         phones=[
             PhoneBenchmarkInfo(name="Pixel4",
                                benchmark_key="5538704950034432",
@@ -159,7 +169,9 @@ def get_s20_default_target_list(batch_config=None):
         ]),
     ModelBenchmarkInfo(
         name="mobilenet-v2",
+        model_artifacts_name="mobilenet-v2.tar.gz",
         model_path="mobilenet-v2/iree_input.mlir",
+        flagfile_path="mobilenet-v2/flagfile",
         phones=[
             PhoneBenchmarkInfo(name="Pixel4",
                                benchmark_key="6338759231537152",
@@ -175,5 +187,9 @@ def get_s20_default_target_list(batch_config=None):
 ]
 
 
+def get_flagfile_name(model_name):
+  return f"{model_name}_flagfile"
+
+
 def get_module_name(model_name, phone_name, mako_tag):
   return f"{model_name}_{phone_name}_{mako_tag}.vmfb"

build_tools/mako/prepare_benchmark_files.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Prepares files for
+#   * compile_android_modules.py
+#   * benchmark_modules_on_android.py
+#
+# The script assumes model artifacts are in google bucket and they are zipped in
+# tar.gz format.
+
+import subprocess
+
+import configuration
+
+
+def main() -> None:
+  for model_benchmark in configuration.MODEL_BENCHMARKS:
+    print(f"Preparing benchmark files for {model_benchmark.name}")
+    subprocess.run(args=[
+        "gsutil", "cp",
+        f"gs://iree-model-artifacts/{model_benchmark.model_artifacts_name}", "."
+    ],
+                   check=True)
+    subprocess.run(args=["tar", "-xvf", model_benchmark.model_artifacts_name],
+                   check=True)
+    subprocess.run(args=[
+        "cp", model_benchmark.flagfile_path,
+        configuration.get_flagfile_name(model_benchmark.name)
+    ],
+                   check=True)
+
+
+if __name__ == "__main__":
+  main()