TensorFlow XLA support (#559)

itlab-vision · Dec 1, 2024 · d944b0b · d944b0b
1 parent 2976a9d
commit d944b0b
Show file tree

Hide file tree

Showing 11 changed files with 34 additions and 3 deletions.
diff --git a/src/benchmark/frameworks/config_parser/test_reporter.py b/src/benchmark/frameworks/config_parser/test_reporter.py
@@ -65,6 +65,7 @@ def prepare_framework_params(self):
         match_parameter_description['kmp_affinity'] = 'KMP_AFFINITY'
         match_parameter_description['num_inter_threads'] = 'Inter threads'
         match_parameter_description['num_intra_threads'] = 'Intra threads'
+        match_parameter_description['use_xla'] = 'Using XLA'
 
         match_parameter_description['execution_providers'] = 'Execution Providers'
         match_parameter_description['execution_mode'] = 'Execution Mode'

diff --git a/src/benchmark/frameworks/tensorflow/tensorflow_parameters_parser.py b/src/benchmark/frameworks/tensorflow/tensorflow_parameters_parser.py
@@ -15,6 +15,7 @@ def parse_parameters(self, curr_test):
         CONFIG_FRAMEWORK_DEPENDENT_INTER_OP_PARALLELISM_THREADS_TAG = 'InterOpParallelismThreads'
         CONFIG_FRAMEWORK_DEPENDENT_INTRA_OP_PARALLELISM_THREADS_TAG = 'IntraOpParallelismThreads'
         CONFIG_FRAMEWORK_DEPENDENT_KMP_AFFINITY_TAG = 'KmpAffinity'
+        CONFIG_FRAMEWORK_DEPENDENT_USE_XLA_TAG = 'UseXLA'
 
         dep_parameters_tag = curr_test.getElementsByTagName(CONFIG_FRAMEWORK_DEPENDENT_TAG)[0]
 
@@ -38,6 +39,8 @@ def parse_parameters(self, curr_test):
             CONFIG_FRAMEWORK_DEPENDENT_INTRA_OP_PARALLELISM_THREADS_TAG)[0].firstChild
         _kmp_affinity = dep_parameters_tag.getElementsByTagName(
             CONFIG_FRAMEWORK_DEPENDENT_KMP_AFFINITY_TAG)[0].firstChild
+        _use_xla = dep_parameters_tag.getElementsByTagName(
+            CONFIG_FRAMEWORK_DEPENDENT_USE_XLA_TAG)[0].firstChild
 
         return TensorFlowParameters(
             channel_swap=_channel_swap.data if _channel_swap else None,
@@ -50,12 +53,13 @@ def parse_parameters(self, curr_test):
             inter_op_parallelism_threads=_inter_op_parallelism_threads.data if _inter_op_parallelism_threads else None,
             intra_op_parallelism_threads=_intra_op_parallelism_threads.data if _intra_op_parallelism_threads else None,
             kmp_affinity=_kmp_affinity.data if _kmp_affinity else None,
+            use_xla=_use_xla.data if _use_xla else None,
         )
 
 
 class TensorFlowParameters(FrameworkParameters):
     def __init__(self, channel_swap, mean, input_scale, input_shape, input_name, output_names, thread_count,
-                 inter_op_parallelism_threads, intra_op_parallelism_threads, kmp_affinity):
+                 inter_op_parallelism_threads, intra_op_parallelism_threads, kmp_affinity, use_xla):
         self.channel_swap = None
         self.mean = None
         self.input_scale = None
@@ -66,6 +70,7 @@ def __init__(self, channel_swap, mean, input_scale, input_shape, input_name, out
         self.num_inter_threads = None
         self.num_intra_threads = None
         self.kmp_affinity = None
+        self.use_xla = None
 
         if self._parameter_is_not_none(channel_swap):
             if self._channel_swap_is_correct(channel_swap):
@@ -108,6 +113,8 @@ def __init__(self, channel_swap, mean, input_scale, input_shape, input_name, out
                 raise ValueError('Intra op parallelism threads can only take integer value')
         if self._parameter_is_not_none(kmp_affinity):
             self.kmp_affinity = kmp_affinity
+        if self._parameter_is_not_none(use_xla):
+            self.use_xla = use_xla
 
     def _input_shape_is_correct(self, input_shape):
         shape_check = input_shape.split()

diff --git a/src/benchmark/frameworks/tensorflow/tensorflow_process.py b/src/benchmark/frameworks/tensorflow/tensorflow_process.py
@@ -68,13 +68,21 @@ def _fill_command_line(self):
             common_params = self._add_argument_to_cmd_line(common_params, '--raw_output', 'true')
         common_params = self._add_flag_to_cmd_line(common_params, '--restrisct_gpu_usage')
 
+        use_xla = self._test.dep_parameters.use_xla
+        if use_xla:
+            common_params = self._add_optional_argument_to_cmd_line(common_params, '--use_xla', 'true')
+
         command_line = f'{python} {path_to_tensorflow_script} {common_params}'
 
         nthreads = self._test.dep_parameters.nthreads
         if nthreads:
             command_line = self._add_env_to_cmd_line(command_line, 'OMP_NUM_THREADS', nthreads)
         kmp_affinity = self._test.dep_parameters.kmp_affinity
         if kmp_affinity:
-            command_line = TensorFlowProcess._add_env_to_cmd_line(command_line, 'KMP_AFFINITY', kmp_affinity)
+            command_line = self._add_env_to_cmd_line(command_line, 'KMP_AFFINITY', kmp_affinity)
+        if use_xla:
+            command_line = self._add_env_to_cmd_line(command_line, 'TF_XLA_FLAGS',
+                                                     '"--tf_xla_auto_jit=2 \
+                                                     --tf_xla_cpu_global_jit"')
 
         return command_line
diff --git a/src/configs/README.md b/src/configs/README.md
@@ -140,6 +140,7 @@
   - `KmpAffinity` - опциональный тег. Позволяет установить значение переменной
     окружения KMP_AFFINITY. По умолчанию не задан. Подробнее про атрибуты, принимаемые
     переменной окружения, [здесь][kmp-affinity-docs].
+  - `UseXLA` - опциональный тег. Включает поддержку XLA-компилятора для запуска моделей. По умолчанию False.
 
 - Набор тегов для тестирования вывода средствами ONNX Runtime CPP,
   OpenCV DNN CPP, PyTorch CPP and TensorFlow Lite CPP:

diff --git a/src/configs/benchmark_configuration_file_template.xml b/src/configs/benchmark_configuration_file_template.xml
@@ -97,6 +97,7 @@
             <InterOpParallelismThreads></InterOpParallelismThreads>
             <IntraOpParallelismThreads></IntraOpParallelismThreads>
             <KmpAffinity></KmpAffinity>
+            <UseXLA></UseXLA>
         </FrameworkDependent>
     </Test>
     <Test>

diff --git a/src/inference/README.md b/src/inference/README.md
@@ -363,6 +363,7 @@ inference_tensorflow.py
   TensorFlow автоматически.
 - `--raw_output` - работа скрипта без логов. По умолчанию не установлен.
 - `--restrisct_gpu_usage` - ограничение использования видеокарт до 1-й. По умолчанию не установлен.
+- `--use_xla` - включить использование XLA-компилятора. По умолчанию `False`.
 
 #### Примеры запуска
 

diff --git a/src/inference/inference_tensorflow.py b/src/inference/inference_tensorflow.py
@@ -144,6 +144,12 @@ def cli_argument_parser():
     parser.add_argument('--restrisct_gpu_usage',
                         action='store_true',
                         help='Restrict TensorFlow to only use the first GPU')
+    parser.add_argument('--use_xla',
+                        help='Enable XLA-compiler',
+                        required=False,
+                        default=False,
+                        type=bool,
+                        dest='use_xla')
     args = parser.parse_args()
 
     return args
@@ -248,7 +254,6 @@ def main():
     report_writer.update_configuration_setup(batch_size=args.batch_size,
                                              iterations_num=args.number_iter,
                                              target_device=args.device)
-
     if args.device == 'NVIDIA_GPU':
         if is_gpu_available():
             if args.restrisct_gpu_usage:
@@ -260,6 +265,9 @@ def main():
     if args.device == 'CPU' and is_gpu_available():
         log.warning(f'NVIDIA_GPU device(s) {get_gpu_devices()} available on machine,'
                     f' tensorflow will use NVIDIA_GPU by default')
+    if args.use_xla:
+        tf.config.optimizer.set_jit(True)
+    log.info(f'Current XLA status: {tf.config.optimizer.get_jit()}')
 
     input_name = args.input_name
     input_op_name = get_input_operation_name(input_name)

diff --git a/tests/smoke_test/configs/dl_models/mobilenet-v1-1.0-224-tf_TensorFlow.xml b/tests/smoke_test/configs/dl_models/mobilenet-v1-1.0-224-tf_TensorFlow.xml
@@ -31,6 +31,7 @@
             <InterOpParallelismThreads></InterOpParallelismThreads>
             <IntraOpParallelismThreads></IntraOpParallelismThreads>
             <KmpAffinity></KmpAffinity>
+            <UseXLA></UseXLA>
         </FrameworkDependent>
     </Test>
 </Tests>
diff --git a/tests/smoke_test/configs/dl_models/mobilenet-v1-1.0-224-tf_TensorFlow_classification.xml b/tests/smoke_test/configs/dl_models/mobilenet-v1-1.0-224-tf_TensorFlow_classification.xml
@@ -32,6 +32,7 @@
             <InterOpParallelismThreads></InterOpParallelismThreads>
             <IntraOpParallelismThreads></IntraOpParallelismThreads>
             <KmpAffinity></KmpAffinity>
+            <UseXLA></UseXLA>
         </FrameworkDependent>
     </Test>
 </Tests>
diff --git a/tests/smoke_test/configs/dl_models/yolo-v2-tiny-tf-TensorFlow_yolo_v2_tiny_coco.xml b/tests/smoke_test/configs/dl_models/yolo-v2-tiny-tf-TensorFlow_yolo_v2_tiny_coco.xml
@@ -32,6 +32,7 @@
             <InterOpParallelismThreads></InterOpParallelismThreads>
             <IntraOpParallelismThreads></IntraOpParallelismThreads>
             <KmpAffinity></KmpAffinity>
+            <UseXLA></UseXLA>
         </FrameworkDependent>
     </Test>
 </Tests>
diff --git a/tests/smoke_test/configs/dl_models/yolo-v3-tf-TensorFlow_yolo_v3_tf.xml b/tests/smoke_test/configs/dl_models/yolo-v3-tf-TensorFlow_yolo_v3_tf.xml
@@ -32,6 +32,7 @@
             <InterOpParallelismThreads></InterOpParallelismThreads>
             <IntraOpParallelismThreads></IntraOpParallelismThreads>
             <KmpAffinity></KmpAffinity>
+            <UseXLA></UseXLA>
         </FrameworkDependent>
     </Test>
 </Tests>