Only generate enabled models

iree-org · Sep 12, 2023 · 0da1480 · 0da1480
1 parent 0bb7300
commit 0da1480
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 7 deletions.
diff --git a/.github/workflows/update_model_artifacts.yml b/.github/workflows/update_model_artifacts.yml
@@ -63,12 +63,7 @@ jobs:
             --env "AUTO_UPLOAD=0" \
             --env "OUTPUT_DIR=jax" \
             "gcr.io/iree-oss/openxla-benchmark/cuda11.8-cudnn8.9@sha256:f43984cd6c16ad1faad4dfb6aac3f53e552dd728c9330c90752e78ae51e4276f" \
-            "common_benchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/generate_model_artifacts.sh" \
-              "T5_LARGE_FP32_JAX_512XI32_BATCH1" \
-              "T5_LARGE_FP32_JAX_512XI32_BATCH16" \
-              "T5_LARGE_FP32_JAX_512XI32_BATCH24" \
-              "T5_LARGE_FP32_JAX_512XI32_BATCH32" \
-              "T5_LARGE_FP32_JAX_512XI32_BATCH48"
+            "common_benchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/generate_model_artifacts.sh"
       - name: "Upload JAX model artifacts"
         run: |
-          gcloud storage cp -r "jax/*"  "gs://iree-model-artifacts/jax/"
+          gcloud storage cp -r "jax/*"  "gs://iree-model-artifacts/jax"
diff --git a/...nchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/generate_model_artifacts.sh b/...nchmark_suite/openxla/benchmark/comparative_suite/jax/scripts/generate_model_artifacts.sh
@@ -59,6 +59,16 @@ declare -a args=(
 
 if (( "${#FILTER[@]}" > 0 )); then
   args+=( --filter "${FILTER[@]}" )
+else
+  # Generate enabled models in comparative_benchmark/jax/benchmark_xla.sh
+  args+=(
+    --filter
+    "RESNET50_FP32_JAX_.+"
+    "BERT_LARGE_FP32_JAX_.+_BATCH(1|16|24|32|48|64|512)"
+    "T5_LARGE_FP32_JAX_.+_BATCH(1|16|24|32|48|64)"
+    "T5_4CG_LARGE_FP32_JAX_.+"
+    "GPT2LMHEAD_FP32_JAX_.+"
+  )
 fi
 
 if (( AUTO_UPLOAD == 1 )); then