--use_xla_auto_jit added to the benchmark

Author: DEKHTIARJonathan

tftrt/examples/benchmark_args.py

@@ -140,6 +140,14 @@ def __init__(self):
             help="If set to True, the benchmark will use XLA JIT Compilation."
         )
 
+        self._add_bool_argument(
+            name="use_xla_auto_jit",
+            default=False,
+            required=False,
+            help="If set to True, the benchmark will use XLA JIT Auto "
+            "Clustering Compilation."
+        )
+
         self._add_bool_argument(
             name="use_synthetic_data",
             default=False,

tftrt/examples/benchmark_runner.py

@@ -57,6 +57,10 @@ def evaluate_model(self, predictions, expected, bypass_data_to_eval):
     def __init__(self, args):
         self._args = args
 
+        if args.use_xla_auto_jit:
+            print("[Benchmark] - Activating XLA JIT Auto Clustering")
+            os.environ["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit"
+
         logging.getLogger("tensorflow").setLevel(logging.INFO)
         logging.disable(logging.WARNING)
 

tftrt/examples/image_classification/base_run_inference.sh

@@ -11,7 +11,6 @@ MODEL_DIR=""
 NVIDIA_TF32_OVERRIDE=""
 
 BYPASS_ARGUMENTS=""
-TF_AUTO_JIT_XLA_FLAG=""
 
 # Loop through arguments and process them
 for arg in "$@"
@@ -39,10 +38,6 @@ do
         --output_tensors_name=*)
         shift # Remove --output_tensors_name= from processing
         ;;
-        --use_xla_auto_jit)
-        TF_AUTO_JIT_XLA_FLAG="TF_XLA_FLAGS=\"--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit\""
-        shift # Remove --use_xla_auto_jit from processing
-        ;;
         *)
         BYPASS_ARGUMENTS=" ${BYPASS_ARGUMENTS} ${arg}"
         ;;
@@ -114,7 +109,6 @@ echo "[*] NUM_CLASSES: ${NUM_CLASSES}"
 echo "[*] MAX_SAMPLES: ${MAX_SAMPLES}"
 echo "[*] OUTPUT_TENSORS_NAME: ${OUTPUT_TENSORS_NAME}"
 echo ""
-echo "[*] TF_AUTO_JIT_XLA_FLAG: ${TF_AUTO_JIT_XLA_FLAG}"
 echo "[*] BYPASS_ARGUMENTS: $(echo \"${BYPASS_ARGUMENTS}\" | tr -s ' ')"
 echo -e "********************************************************************\n"
 
@@ -158,7 +152,7 @@ cd ${BENCH_DIR}
 
 # Execute the example
 
-PREPEND_COMMAND="${TF_AUTO_JIT_XLA_FLAG} ${NVIDIA_TF32_OVERRIDE}"
+PREPEND_COMMAND="${NVIDIA_TF32_OVERRIDE}"
 
 COMMAND="${PREPEND_COMMAND} python image_classification.py \
     --data_dir ${DATA_DIR} \

tftrt/examples/object_detection/base_run_inference.sh

@@ -11,7 +11,6 @@ MODEL_DIR=""
 NVIDIA_TF32_OVERRIDE=""
 
 BYPASS_ARGUMENTS=""
-TF_AUTO_JIT_XLA_FLAG=""
 BATCH_SIZE=8
 
 # Loop through arguments and process them
@@ -44,10 +43,6 @@ do
         MODEL_DIR="${arg#*=}"
         shift # Remove --input_saved_model_dir= from processing
         ;;
-        --use_xla_auto_jit)
-        TF_AUTO_JIT_XLA_FLAG="TF_XLA_FLAGS=\"--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit\""
-        shift # Remove --use_xla_auto_jit from processing
-        ;;
         *)
         BYPASS_ARGUMENTS=" ${BYPASS_ARGUMENTS} ${arg}"
         ;;
@@ -84,7 +79,6 @@ echo "[*] MAX_WORKSPACE_SIZE: ${MAX_WORKSPACE_SIZE}"
 echo "[*] MAX_SAMPLES: ${MAX_SAMPLES}"
 echo "[*] OUTPUT_TENSORS_NAME: ${OUTPUT_TENSORS_NAME}"
 echo ""
-echo "[*] TF_AUTO_JIT_XLA_FLAG: ${TF_AUTO_JIT_XLA_FLAG}"
 echo "[*] BYPASS_ARGUMENTS: $(echo \"${BYPASS_ARGUMENTS}\" | tr -s ' ')"
 echo -e "********************************************************************\n"
 
@@ -149,7 +143,7 @@ fi
 
 # Step 2: Execute the example
 
-PREPEND_COMMAND="${TF_AUTO_JIT_XLA_FLAG} ${NVIDIA_TF32_OVERRIDE}"
+PREPEND_COMMAND="${NVIDIA_TF32_OVERRIDE}"
 
 COMMAND="${PREPEND_COMMAND} python object_detection.py \
     --data_dir ${VAL_DATA_DIR} \

tftrt/examples/transformers/base_run_inference.sh

@@ -13,7 +13,6 @@ NVIDIA_TF32_OVERRIDE=""
 DATA_DIR="/tmp"
 
 BYPASS_ARGUMENTS=""
-TF_AUTO_JIT_XLA_FLAG=""
 BATCH_SIZE=32
 SEQ_LEN=128
 
@@ -46,10 +45,6 @@ do
         MODEL_DIR="${arg#*=}"
         shift # Remove --input_saved_model_dir= from processing
         ;;
-        --use_xla_auto_jit)
-        TF_AUTO_JIT_XLA_FLAG="TF_XLA_FLAGS=\"--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit\""
-        shift # Remove --use_xla_auto_jit from processing
-        ;;
         --vocab_size=*)
         shift # Remove --vocab_size= from processing
         ;;
@@ -104,7 +99,6 @@ echo "[*] MAX_WORKSPACE_SIZE: ${MAX_WORKSPACE_SIZE}"
 echo "[*] MAX_SAMPLES: ${MAX_SAMPLES}"
 echo "[*] OUTPUT_TENSORS_NAME: ${OUTPUT_TENSORS_NAME}"
 echo ""
-echo "[*] TF_AUTO_JIT_XLA_FLAG: ${TF_AUTO_JIT_XLA_FLAG}"
 echo "[*] BYPASS_ARGUMENTS: $(echo \"${BYPASS_ARGUMENTS}\" | tr -s ' ')"
 
 echo -e "********************************************************************\n"
@@ -149,7 +143,7 @@ cd ${BENCH_DIR}
 
 # Execute the example
 
-PREPEND_COMMAND="${TF_AUTO_JIT_XLA_FLAG} ${NVIDIA_TF32_OVERRIDE}"
+PREPEND_COMMAND="${NVIDIA_TF32_OVERRIDE}"
 
 COMMAND="${PREPEND_COMMAND} python transformers.py \
     --data_dir ${DATA_DIR} \