[madlib] 02/03: DL: Enable JIT XLA auto-clustering, if available.

[do...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

domino pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/madlib.git

commit 229d2b1cac9b19cee298f96ebf3c28818a0cd038
Author: Domino Valdano <dv...@vmware.com>
AuthorDate: Wed Nov 11 23:22:25 2020 -0800

    DL: Enable JIT XLA auto-clustering, if available.
    
    This is a simple optimization for CPU and/or GPU we can turn on,
    to improve performance.
    
    If it's not available, that means tensorflow version is too old.
    Will give a warning, suggesting to upgrade to 1.14.0 and continue
    without it.
---
 .../postgres/modules/deep_learning/madlib_keras_wrapper.py_in     | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in b/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
index d7b2d41..7b1b3de 100644
--- a/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
+++ b/src/ports/postgres/modules/deep_learning/madlib_keras_wrapper.py_in
@@ -65,6 +65,13 @@ def reset_cuda_env(value):
         if CUDA_VISIBLE_DEVICES_KEY in os.environ:
             del os.environ[CUDA_VISIBLE_DEVICES_KEY]
 
+def enable_xla():
+    os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit'
+    try:
+        tf.config.optimizer.set_jit(True)
+    except:
+        plpy.warning("This version of tensorflow does not fully support XLA auto-cluster JIT optimization.  HINT:  upgrading to tensorflow 1.14.0 may improve performance.")
+
 def get_device_name_and_set_cuda_env(gpu_count, seg):
     if gpu_count > 0:
         device_name = '/gpu:0'
@@ -90,6 +97,7 @@ def get_keras_session(device_name, gpu_count, segments_per_host):
         config.gpu_options.allow_growth = False
         config.gpu_options.per_process_gpu_memory_fraction = memory_fraction
     session = tf.Session(config=config)
+    enable_xla()
     return session
 
 def clear_keras_session(sess = None):