fix build with deeprec-gpu

scripts/python/common_setup.py

@@ -374,14 +374,15 @@ def update_cpu_specific_setting(args):
 
 def get_tf_info(python_executable):
     output = subprocess.check_output(
-        '{} -c "import tensorflow as tf; print(tf.__version__); print(\'\\n\'.join(tf.sysconfig.get_compile_flags())); print(\'\\n\'.join(tf.sysconfig.get_link_flags()))"'.format(
+        '{} -c "import tensorflow as tf; print(tf.__version__); print(tf.__git_version__); print(\'\\n\'.join(tf.sysconfig.get_compile_flags())); print(\'\\n\'.join(tf.sysconfig.get_link_flags()))"'.format(
             python_executable
         ),
         shell=True,
     ).decode()
     lines = output.split("\n")
     major, minor, _ = lines[0].split(".")  # lines[0] is version like 1.15.0
     is_pai = "PAI" in lines[0]
+    is_deeprec = 'deeprec' in lines[1]
     header_dir, lib_dir, lib_name, cxx11_abi = '', '', '', ''
     for line in lines[1:]:
         if line.startswith("-I"):
@@ -410,7 +411,7 @@ def get_tf_info(python_executable):
             raise Exception("Can not find tensorflow's built-in pb version!")
     else:
         raise Exception("Can not find {PB_HEADER_FILE} in tf's include dir!")
-    return major, minor, is_pai, header_dir, lib_dir, lib_name, cxx11_abi, tf_pb_version
+    return major, minor, is_pai, header_dir, lib_dir, lib_name, cxx11_abi, tf_pb_version, is_deeprec
 
 def deduce_cuda_info():
     """Deduce cuda major and minor version and cuda directory."""

scripts/python/tao_build.py

@@ -201,9 +201,12 @@ def _write(line, cmd="build"):
             tf_lib_name,
             tf_cxx11_abi,
             tf_pb_version,
+            is_deeprec
         ) = get_tf_info(python_bin)
         _opt("cxxopt", f"-D_GLIBCXX_USE_CXX11_ABI={tf_cxx11_abi}")
         _opt("host_cxxopt", f"-D_GLIBCXX_USE_CXX11_ABI={tf_cxx11_abi}")
+        _opt("cxxopt", f"-DTF_IS_DEEPREC={is_deeprec}")
+        _opt("host_cxxopt", f"-DTF_IS_DEEPREC={is_deeprec}")
         _action_env("BLADE_WITH_TF", "1")
         _action_env("IF_CXX11_ABI", int(tf_cxx11_abi))
         _action_env("TF_IS_PAI", int(is_pai))

tao_compiler/mlir/xla/ral/context/stream_executor_based_impl.cc

@@ -1392,6 +1392,9 @@ Status RunCudnnConvolution<int8_t>(
     se::DeviceMemoryBase& result_buffer,
     se::ScratchAllocator* scratch_allocator, se::Stream* stream,
     se::dnn::ProfileResult* profile_result) {
+#if defined(TF_IS_DEEPREC)
+  return errors::Internal("Not supported for DeepREC yet");
+#else
   ConvolutionKind kind = params.kind;
   DeviceMemory<int8_t> input_buf(operand_buffers[0]);
   DeviceMemory<int8_t> filter_buf(operand_buffers[1]);
@@ -1466,6 +1469,7 @@ Status RunCudnnConvolution<int8_t>(
     return errors::Internal("Unable to launch convolution");
   }
   return Status::OK();
+#endif  // TF_IS_DEEPREC
 }
 
 template <typename T>

tensorflow_blade/build.py

@@ -150,7 +150,10 @@ def _config(cfg_name, cmd="build"):
             tf_lib_name,
             tf_cxx11_abi,
             tf_pb_version,
+            is_deeprec
         ) = get_tf_info(which("python3"))
+        _opt("cxxopt", f"-DTF_IS_DEEPREC={is_deeprec}")
+        _opt("host_cxxopt", f"-DTF_IS_DEEPREC={is_deeprec}")
         _action_env("BLADE_WITH_TF", "1")
         _config("cxx11abi_" + tf_cxx11_abi)
         _action_env("IF_CXX11_ABI", int(tf_cxx11_abi))
@@ -260,7 +263,7 @@ def package(args):
 def test(args):
     execute("python3 setup.py cpp_test")
     device_mark = "not gpu_only" if args.device != "gpu" else "not cpu_only"
-    tf_major, _, _, _, _, _, _, _ = get_tf_info(which("python3"))
+    tf_major, _, _, _, _, _, _, _, _ = get_tf_info(which("python3"))
     tf_mark = "not tf1_only" if tf_major == "2" else "not tf2_only"
     execute(f"pytest  tests/ -m '{device_mark} and {tf_mark}' -v --forked")