refactor for internal serving build (#961)

* refactor for internal serving build

* fix non-internal_serving build

* fix internal ci failure

* revert blade gemm deps for bridge

* fix transpose bridge build

* fix

.gitignore

@@ -22,6 +22,7 @@ tao/third_party/blade_service_common
 tao/build.txt
 tao/*.tgz
 third_party/bazel/blade_service_common/*.BUILD
+third_party/bazel/blade_service_common/*.BUILD.tpl
 third_party/bazel/blade_service_common/blade_service_common_workspace.bzl.tpl
 **/*.so
 **/*.so.[0-9]

VERSION

@@ -1 +1 @@
-0.21.05
+0.23.01

scripts/python/common_setup.py

@@ -619,9 +619,6 @@ def symlink_disc_files(args):
         )
     )
 
-    logger.info("linking PatineClient")
-    link_dirs(os.path.join(dir_tf_community, 'tao', 'third_party', 'PatineClient'),
-            os.path.join(dir_platform_alibaba, 'third_party', 'PatineClient'))
     logger.info("linking blade_gemm")
     link_dirs(os.path.join(get_source_root_dir(), 'tf_community', 'tao', 'blade_gemm'),
             os.path.join(dir_platform_alibaba, 'blade_gemm'))

tao/BUILD

@@ -10,6 +10,7 @@ load("@local_config_blade_disc_helper//:build_defs.bzl",
     "disc_build_ip",
     "disc_build_time",
     "if_platform_alibaba",
+    "if_internal_serving",
 )
 load("@org_tensorflow//tensorflow:tensorflow.bzl", "tf_cc_shared_object")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_dcu", "if_rocm")
@@ -102,20 +103,33 @@ cc_library(
     ),
     deps = [
         ":version_header",
-        "@local_config_tf//:tf_header_lib",
-        "@local_config_tf//:libtensorflow_framework",
         "@org_tao_compiler//tensorflow/compiler/mlir/xla/ral:ral_bridge",
     ] + if_platform_alibaba([
         "@blade_service_common//blade_service_common:blade_service_common_deps",
-    ]),
+    ]) + if_internal_serving(
+        [
+            "@org_tensorflow//tensorflow/core:lib",
+            "@org_tensorflow//tensorflow/core:framework",
+            "@org_tensorflow//tensorflow/core:core_cpu",
+            "@org_tensorflow//tensorflow/core:stream_executor_no_cuda",
+            "@org_tensorflow//tensorflow/stream_executor:stream_executor_impl",
+        ],
+        [
+            "@local_config_tf//:tf_header_lib",
+            "@local_config_tf//:libtensorflow_framework",
+        ]
+    ),
     copts = if_cuda([
         "-DGOOGLE_CUDA=1"
     ]) + if_dcu([
         "-DTENSORFLOW_USE_DCU=1"
     ]) + if_rocm([
         "-DTENSORFLOW_USE_ROCM=1",
         "-x rocm",
-    ])+ tf_copts(),
+    ]) + if_internal_serving([
+        "-DTF_1_12",
+        "-DIS_PAI_TF",
+    ]) + tf_copts(),
     linkopts = [
         "-lm",
     ],

tao_compiler/decoupling/BUILD

@@ -1,7 +1,7 @@
 licenses(["notice"])  # Apache 2.0
 
 load("//tensorflow:tensorflow.bzl", "tf_cc_binary", "tf_cc_test")
-load("//tensorflow/compiler/mlir/disc:disc.bzl", "if_platform_alibaba")
+load("//tensorflow/compiler/mlir/util:util.bzl", "if_platform_alibaba")
 
 package_group(
     name = "internal",
@@ -25,7 +25,7 @@ package(
 
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured")
-load("//tensorflow/compiler/mlir/disc:disc.bzl", "if_cuda_or_rocm")
+load("//tensorflow/compiler/mlir/util:util.bzl", "if_cuda_or_rocm")
 load(
     "@com_google_protobuf//:protobuf.bzl",
     "cc_proto_library",

tao_compiler/file_map

@@ -1,5 +1,6 @@
 # relative path under tao_compiler, relative path under tf_community
 decoupling,tensorflow/compiler/decoupling
 mlir/disc,tensorflow/compiler/mlir/disc
+mlir/util,tensorflow/compiler/mlir/util
 mlir/xla/ral,tensorflow/compiler/mlir/xla/ral
 .bazelrc.user,tensorflow/../.bazelrc.user

tao_compiler/mlir/disc/BUILD

@@ -18,52 +18,12 @@ load("@org_tensorflow//tensorflow/core/platform:build_config_root.bzl", "if_stat
 load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "gentbl_filegroup", "td_library")
 load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured")
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured")
-load("//tensorflow/compiler/mlir/disc:disc.bzl",
+load("//tensorflow/compiler/mlir/util:util.bzl",
      "disc_cc_library",
      "if_cuda_or_rocm",
      "if_skip_compute_intensive_fusion",
 )
-load("//tensorflow/compiler/mlir/disc:disc.bzl", "if_torch_disc")
-
-config_setting(
-    name = "is_platform_alibaba",
-    define_values = {"is_platform_alibaba": "true"},
-)
-
-config_setting(
-    name = "is_patine",
-    define_values = {"is_patine": "true"},
-)
-
-config_setting(
-    name = "is_blade_gemm",
-    define_values = {"is_blade_gemm": "true"},
-)
-
-config_setting(
-    name = "is_mkldnn",
-    define_values = {"is_mkldnn": "true"},
-)
-
-config_setting(
-    name = "disc_aarch64",
-    define_values = {"disc_aarch64": "true"},
-)
-
-config_setting(
-    name = "disc_x86",
-    define_values = {"disc_x86": "true"},
-)
-
-config_setting(
-    name = "is_torch_disc",
-    define_values = {"is_torch_disc": "true"},
-)
-
-config_setting(
-    name = "skip_compute_intensive_fusion",
-    define_values = {"skip_compute_intensive_fusion": "true"},
-)
+load("//tensorflow/compiler/mlir/util:util.bzl", "if_torch_disc")
 
 package(
     default_visibility = [":friends"],

tao_compiler/mlir/disc/tests/BUILD

@@ -3,7 +3,7 @@ load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured",)
 load("@local_config_rocm//rocm:build_defs.bzl", "if_rocm_is_configured",)
 load("//tensorflow/core/platform:build_config.bzl", "tf_proto_library",)
 load("//tensorflow/core/platform:build_config_root.bzl", "if_static")
-load("//tensorflow/compiler/mlir/disc:disc.bzl", "disc_cc_library")
+load("//tensorflow/compiler/mlir/util:util.bzl", "disc_cc_library")
 
 # buildifier: disable=same-origin-load
 load("//tensorflow:tensorflow.bzl", "get_compatible_with_cloud")

tao_compiler/mlir/disc/tools/disc-replay/BUILD

@@ -4,7 +4,7 @@ load(
   "tf_cc_test",
 )
 
-load("//tensorflow/compiler/mlir/disc:disc.bzl",
+load("//tensorflow/compiler/mlir/util:util.bzl",
      "disc_cc_library",
      "if_cuda_or_rocm",
 )

tao_compiler/mlir/util/BUILD

@@ -0,0 +1,44 @@
+config_setting(
+    name = "is_platform_alibaba",
+    define_values = {"is_platform_alibaba": "true"},
+)
+
+config_setting(
+    name = "is_patine",
+    define_values = {"is_patine": "true"},
+)
+
+config_setting(
+    name = "is_blade_gemm",
+    define_values = {"is_blade_gemm": "true"},
+)
+
+config_setting(
+    name = "is_mkldnn",
+    define_values = {"is_mkldnn": "true"},
+)
+
+config_setting(
+    name = "disc_aarch64",
+    define_values = {"disc_aarch64": "true"},
+)
+
+config_setting(
+    name = "disc_x86",
+    define_values = {"disc_x86": "true"},
+)
+
+config_setting(
+    name = "is_torch_disc",
+    define_values = {"is_torch_disc": "true"},
+)
+
+config_setting(
+    name = "skip_compute_intensive_fusion",
+    define_values = {"skip_compute_intensive_fusion": "true"},
+)
+
+config_setting(
+    name = "is_internal_serving",
+    define_values = {"is_internal_serving": "true"},
+)

tao_compiler/mlir/disc/disc.bzl → tao_compiler/mlir/util/util.bzl

@@ -22,55 +22,61 @@ def disc_cc_library(copts = tf_copts(), **kwargs):
 
 def if_platform_alibaba(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:is_platform_alibaba": if_true,
+        "//tensorflow/compiler/mlir/util:is_platform_alibaba": if_true,
         "//conditions:default": if_false
     })
 
 def if_cuda_or_rocm(if_true, if_false=[]):
     return select({
-        "@local_config_cuda//:is_cuda_enabled": if_true,
+        "@local_config_cuda//cuda:using_nvcc": if_true,
         "@local_config_rocm//rocm:using_hipcc": if_true,
         "//conditions:default": if_false
     })
 
 def if_patine(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:is_patine": if_true,
+        "//tensorflow/compiler/mlir/util:is_patine": if_true,
         "//conditions:default": if_false
     })
 
 def if_blade_gemm(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:is_blade_gemm": if_true,
+        "//tensorflow/compiler/mlir/util:is_blade_gemm": if_true,
         "//conditions:default": if_false
     })
 
 def if_mkldnn(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:is_mkldnn": if_true,
+        "//tensorflow/compiler/mlir/util:is_mkldnn": if_true,
         "//conditions:default": if_false
     })
 
 def if_disc_aarch64(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:disc_aarch64": if_true,
+        "//tensorflow/compiler/mlir/util:disc_aarch64": if_true,
         "//conditions:default": if_false
     })
 
 def if_disc_x86(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:disc_x86": if_true,
+        "//tensorflow/compiler/mlir/util:disc_x86": if_true,
         "//conditions:default": if_false
     })
 
 def if_torch_disc(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:is_torch_disc": if_true,
+        "//tensorflow/compiler/mlir/util:is_torch_disc": if_true,
         "//conditions:default": if_false
     })
 
 def if_skip_compute_intensive_fusion(if_true, if_false=[]):
     return select({
-        "//tensorflow/compiler/mlir/disc:skip_compute_intensive_fusion": if_true,
+        "//tensorflow/compiler/mlir/util:skip_compute_intensive_fusion": if_true,
+        "//conditions:default": if_false
+    })
+
+def if_internal_serving(if_true, if_false=[]):
+    return select({
+        "@local_config_blade_disc_helper//:is_internal_serving": if_true,
         "//conditions:default": if_false
     })