[Feat] Rebase intel-xla to jax 0.4.24 and upgrade PJRT API from 0.33 to 0.40 (#2691)

cboss6 · web-flow · commit 6bbfe9abc697 · 2024-06-04T08:55:46.000+08:00
diff --git a/.bazelrc b/.bazelrc
@@ -15,7 +15,7 @@ build:cpu --define=device=cpu
 # build:cpu --define=build_with_tbb=true
 
 # This config option is used for GPU backend. (Deprecated)
-build:gpu --crosstool_top=@local_config_sycl//crosstool_sycl:toolchain
+build:gpu --crosstool_top=@itex_local_config_sycl//crosstool_sycl:toolchain
 build:gpu --define=using_sycl=true --define=build_with_sycl=true --define=using_xetla=true --define=using_nextpluggable=true
 build:gpu --define=build_with_onednn_graph=true
 build:gpu --define=device=gpu
@@ -24,15 +24,15 @@ build:gpu --repo_env TF_NEED_SYCL=1
 
 # This config option is used for GPU backend.
 # xpu will replace gpu as default config option for GPU only backend.
-build:xpu --crosstool_top=@local_config_sycl//crosstool_sycl:toolchain
+build:xpu --crosstool_top=@itex_local_config_sycl//crosstool_sycl:toolchain
 build:xpu --define=using_sycl=true --define=build_with_sycl=true --define=using_xetla=true --define=using_nextpluggable=true
 build:xpu --define=build_with_onednn_graph=true
 build:xpu --define=device=gpu
 build:xpu --define=tensorflow_mkldnn_contraction_kernel=0
 build:gpu --repo_env TF_NEED_SYCL=1
 
 # This config option is used for XPU backend (CPU + GPU). (Disabled)
-# build:xpu --crosstool_top=@local_config_sycl//crosstool_sycl:toolchain
+# build:xpu --crosstool_top=@itex_local_config_sycl//crosstool_sycl:toolchain
 # build:xpu --define=using_xetla=true
 # build:xpu --define=build_as_xpu=true
 # build:xpu --define=build_with_onednn_graph=true
diff --git a/WORKSPACE b/WORKSPACE
@@ -32,11 +32,11 @@ load(
 #    curl -L https://github.com/openxla/xla/archive/<git hash>.tar.gz | sha256sum
 #    and update XLA_SHA256 with the result.
 
-XLA_COMMIT = "cf3c9138f9052e7a7761344693b2c677a6161f6d"
+XLA_COMMIT = "98bf412acc77c515ed7b40cff26ac1f1decf8f5b"
 
 new_git_repository(
-    name = "intel_xla",
-    build_file = clean_dep("//third_party/intel_xla:intel_xla.BUILD"),
+    name = "intel_extension_for_openxla",
+    build_file = clean_dep("//third_party/intel_extension_for_openxla:intel_extension_for_openxla.BUILD"),
     commit = XLA_COMMIT,
     remote = "https://github.com/intel/intel-extension-for-openxla.git",
     verbose = True,
@@ -55,10 +55,6 @@ new_git_repository(
 #    path = "/path/to/xla",
 # )
 
-load("@intel_xla//xla:workspace.bzl", "workspace")
-
-workspace()
-
 # To update XLA to a new revision,
 # a) update URL and strip_prefix to the new git commit hash
 # b) get the sha256 hash of the commit by running:
@@ -67,11 +63,11 @@ workspace()
 http_archive(
     name = "xla",
     patch_args = ["-p1"],
-    patches = ["//third_party:openxla.patch"],
-    sha256 = "637acdfce5a406bd929bca052863be8200afbaac0daecbc547e9aff1967d7c00",
-    strip_prefix = "xla-a5a30f7b34b69b7cada6deb4b39f77788e28369c",
+    patches = ["@intel_extension_for_openxla//third_party:openxla.patch"],
+    sha256 = "db007b6628cfe108c63f45d611c6de910abe3ee827e55f08314ce143c4887d66",
+    strip_prefix = "xla-12eee889e1f2ad41e27d7b0e970cb92d282d3ec5",
     urls = [
-        "https://github.com/openxla/xla/archive/a5a30f7b34b69b7cada6deb4b39f77788e28369c.tar.gz",
+        "https://github.com/openxla/xla/archive/12eee889e1f2ad41e27d7b0e970cb92d282d3ec5.tar.gz",
     ],
 )
 
@@ -87,6 +83,10 @@ http_archive(
 #    path = "/path/to/xla",
 # )
 
+load("@intel_extension_for_openxla//xla:workspace.bzl", "workspace")
+
+workspace()
+
 load("@xla//:workspace4.bzl", "xla_workspace4")
 
 xla_workspace4()
diff --git a/itex/BUILD b/itex/BUILD
@@ -2,7 +2,7 @@ package(default_visibility = ["//visibility:public"])
 
 load("//itex:itex.bzl", "cc_binary", "if_using_nextpluggable_device", "itex_xpu_binary", "itex_xpu_library", "tf_copts")
 load("@bazel_skylib//rules:common_settings.bzl", "string_flag")
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 load("@bazel_skylib//lib:selects.bzl", "selects")
 
 string_flag(
@@ -69,7 +69,7 @@ selects.config_setting_group(
 selects.config_setting_group(
     name = "if_using_sycl",
     match_any = [
-        "@local_config_sycl//sycl:using_sycl",
+        "@itex_local_config_sycl//sycl:using_sycl",
         ":gpu_build",
     ],
 )
@@ -78,7 +78,7 @@ selects.config_setting_group(
     name = "gpu_with_nextpluggable_device_build",
     match_all = [
         "@intel_extension_for_tensorflow//itex:gpu_build",
-        "@local_config_sycl//sycl:using_nextpluggable",
+        "@itex_local_config_sycl//sycl:using_nextpluggable",
     ],
 )
 
@@ -94,7 +94,7 @@ cc_binary(
         "//itex/core/wrapper:itex_gpu_wrapper",
     ],
     copts = select({
-        "@local_config_sycl//sycl:using_nextpluggable": ["-DUSING_NEXTPLUGGABLE_DEVICE"],
+        "@itex_local_config_sycl//sycl:using_nextpluggable": ["-DUSING_NEXTPLUGGABLE_DEVICE"],
         "//conditions:default": [],
     }),
     linkopts = ["-Wl,-rpath,$$ORIGIN/../intel_extension_for_tensorflow"],
@@ -131,7 +131,7 @@ itex_xpu_binary(
         "//itex/core/kernels:libitex_common",
         "@local_config_tf//:_pywrap_tensorflow_internal",
     ] + if_using_nextpluggable_device([
-        "//itex/core/utils:libintel_xla",
+        "//itex/core/utils:libintel_extension_for_openxla",
     ]),
 )
 
@@ -152,7 +152,7 @@ itex_xpu_binary(
         "//itex/core:protos_all_cc",
         "//itex/core/kernels:itex_common_cc",
     ] + if_using_nextpluggable_device([
-        "//itex/core/utils:libintel_xla",
+        "//itex/core/utils:libintel_extension_for_openxla",
     ]),
 )
 
diff --git a/itex/core/devices/BUILD b/itex/core/devices/BUILD
@@ -1,7 +1,7 @@
 package(default_visibility = ["//visibility:public"])
 
 load("//itex:itex.bzl", "cc_library", "if_using_nextpluggable_device")
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 
 cc_library(
     name = "xpu_device_util",
diff --git a/itex/core/devices/next_pluggable_device/BUILD b/itex/core/devices/next_pluggable_device/BUILD
@@ -10,7 +10,7 @@ cc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//itex/core/devices/gpu:itex_gpu_runtime_imp",
-        "//itex/core/utils:libintel_xla",
+        "//itex/core/utils:libintel_extension_for_openxla",
         "//third_party/build_option/sycl:itex_gpu_header",
     ],
     alwayslink = True,
diff --git a/itex/core/graph/BUILD b/itex/core/graph/BUILD
@@ -1,5 +1,5 @@
 load("//itex:itex.bzl", "cc_library")
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 
 cc_library(
     name = "xpu_optimizer",
diff --git a/itex/core/kernels/BUILD b/itex/core/kernels/BUILD
@@ -1,5 +1,5 @@
 load("//itex:itex.bzl", "cc_binary", "cc_library", "if_using_nextpluggable_device", "itex_xpu_library", "tf_copts")
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl", "if_xetla")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl", "if_xetla")
 
 XPU_KERNEL_DEPS = [
     "//itex/core/devices:device_backend_util_hdr",
diff --git a/itex/core/kernels/onednn/block/BUILD b/itex/core/kernels/onednn/block/BUILD
@@ -1,5 +1,5 @@
 load("//itex:itex.bzl", "itex_xpu_library", "tf_copts")
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 
 # List all kernels here.
 itex_xpu_library(
diff --git a/itex/core/kernels/onednn_graph/BUILD b/itex/core/kernels/onednn_graph/BUILD
@@ -10,7 +10,7 @@ cc_library(
         "//itex:core",
         "//itex/core/utils/onednn:onednn_graph_util",
     ] + if_using_nextpluggable_device([
-        "//itex/core/utils:libintel_xla",
+        "//itex/core/utils:libintel_extension_for_openxla",
     ]),
     alwayslink = True,
 )
diff --git a/itex/core/utils/BUILD b/itex/core/utils/BUILD
@@ -3,7 +3,7 @@ load(
     "//itex/core/utils:build_config.bzl",
     "tf_protobuf_deps",
 )
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 
 package(
     default_visibility = [
@@ -102,7 +102,7 @@ cc_library(
         "@nsync//:nsync_cpp",
     ] + if_using_nextpluggable_device([
         "//third_party/build_option/sycl:itex_gpu_header",
-        ":libintel_xla",
+        ":libintel_extension_for_openxla",
     ]),
 )
 
@@ -422,18 +422,18 @@ cc_library(
 )
 
 cc_import(
-    name = "libintel_xla",
-    shared_library = ":intel_xla",
+    name = "libintel_extension_for_openxla",
+    shared_library = ":intel_extension_for_openxla",
     visibility = ["//visibility:public"],
 )
 
 itex_xpu_binary(
-    name = "libintel_xla.so",
+    name = "libintel_extension_for_openxla.so",
     linkshared = True,
     set_target = "gpu_orig_backend",
     visibility = ["//visibility:public"],
     deps = [
-        "@intel_xla//xla/pjrt:tf_pjrt_helper",
+        "@intel_extension_for_openxla//xla/pjrt:tf_pjrt_helper",
         "@onednn_gpu",
     ] + select({
         "//third_party/onednn:build_with_onednn_graph": ["@onednn_gpu//:onednn_graph_gpu"],
diff --git a/itex/itex.bzl b/itex/itex.bzl
@@ -1,6 +1,6 @@
 # Return the options to use for a C++ library or binary build.
 # Uses the ":optmode" config_setting to pick the options.
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl", "if_xetla")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl", "if_xetla")
 load("@bazel_skylib//lib:selects.bzl", "selects")
 
 def if_linux_x86_64(a, otherwise = []):
@@ -58,7 +58,7 @@ def if_cc_threadpool_build(if_true, if_false = []):
 
 def if_gpu_backend(if_true, if_false = []):
     return selects.with_or({
-        ("@local_config_sycl//sycl:using_sycl", "@intel_extension_for_tensorflow//itex:xpu_build"): if_true,
+        ("@itex_local_config_sycl//sycl:using_sycl", "@intel_extension_for_tensorflow//itex:xpu_build"): if_true,
         "//conditions:default": if_false,
     })
 
diff --git a/itex/python/build_defs.bzl b/itex/python/build_defs.bzl
@@ -1,6 +1,6 @@
 """Build rules for pybind11"""
 
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 
 PYBIND_COPTS = [
     "-fexceptions",
diff --git a/itex/tools/pip_package/BUILD b/itex/tools/pip_package/BUILD
@@ -15,7 +15,7 @@ load(
     "omp_deps",
 )
 load(
-    "@local_config_sycl//sycl:build_defs.bzl",
+    "@itex_local_config_sycl//sycl:build_defs.bzl",
     "if_sycl",
 )
 load(
diff --git a/itex/tools/pip_package/build_pip_package.sh b/itex/tools/pip_package/build_pip_package.sh
@@ -175,7 +175,7 @@ function prepare_src() {
   cp ${RUNFILES}/../../../../core/kernels/libitex_common.so ${LIB_TMPDIR}/intel_extension_for_tensorflow/
   xetla_lib=$(find ${RUNFILES} -name libitex_gpu_xetla.so)
   [ -f "${xetla_lib}" ] && cp ${xetla_lib} ${LIB_TMPDIR}/intel_extension_for_tensorflow/
-  xla_so_lib=$(find ${RUNFILES} -iname libintel_xla.so)
+  xla_so_lib=$(find ${RUNFILES} -iname libintel_extension_for_openxla.so)
   [ -f "${xla_so_lib}" ] && cp ${xla_so_lib} ${LIB_TMPDIR}/intel_extension_for_tensorflow/
 
   mkdir ${LIB_TMPDIR}/intel_extension_for_tensorflow/python
diff --git a/itex/workspace.bzl b/itex/workspace.bzl
@@ -32,7 +32,7 @@ def _itex_bind():
 
 def itex_workspace(path_prefix = "", tf_repo_name = ""):
     """All external dependencies for TF builds"""
-    sycl_configure(name = "local_config_sycl")
+    sycl_configure(name = "itex_local_config_sycl")
     syslibs_configure(name = "local_config_syslibs")
 
     farmhash()
diff --git a/test/cc/label_image/Makefile.tpl b/test/cc/label_image/Makefile.tpl
@@ -7,7 +7,7 @@ ITEX_CC_PATH= <ITEX_PATH>
 include = -I $(TF_INCLUDE_PATH)
 
 ifeq ($(backend), GPU)
-    lib = -L $(TFCC_PATH) -L $(ITEX_CC_PATH) -lintel_xla -ltensorflow_framework -ltensorflow_cc
+    lib = -L $(TFCC_PATH) -L $(ITEX_CC_PATH) -lintel_extension_for_openxla -ltensorflow_framework -ltensorflow_cc
 else
     lib = -L $(TFCC_PATH) -ltensorflow_framework -ltensorflow_cc
 endif
diff --git a/test/cc/simple_matmul/Makefile.tpl b/test/cc/simple_matmul/Makefile.tpl
@@ -7,7 +7,7 @@ ITEX_CC_PATH= <ITEX_PATH>
 include = -I $(TF_INCLUDE_PATH)
 
 ifeq ($(backend), GPU)
-    lib = -L $(TFCC_PATH) -L $(ITEX_CC_PATH) -lintel_xla -ltensorflow_framework -ltensorflow_cc
+    lib = -L $(TFCC_PATH) -L $(ITEX_CC_PATH) -lintel_extension_for_openxla -ltensorflow_framework -ltensorflow_cc
 else
     lib = -L $(TFCC_PATH) -ltensorflow_framework -ltensorflow_cc
 endif
diff --git a/third_party/absl/com_google_absl.BUILD b/third_party/absl/com_google_absl.BUILD
diff --git a/third_party/absl/workspace.bzl b/third_party/absl/workspace.bzl
@@ -35,16 +35,3 @@ def repo():
         "//third_party/absl:system.absl.{name}.BUILD".format(name = n): "absl/{name}/BUILD.bazel".format(name = n)
         for n in SYS_DIRS
     }
-
-    tf_http_archive(
-        name = "com_google_absl",
-        sha256 = ABSL_SHA256,
-        build_file = "//third_party/absl:com_google_absl.BUILD",
-        system_build_file = "//third_party/absl:system.BUILD",
-        system_link_files = SYS_LINKS,
-        # This patch pulls in a fix for designated initializers that MSVC
-        # complains about. It shouldn't be necessary at the next LTS release.
-        patch_file = ["//third_party/absl:absl_designated_initializers.patch"],
-        strip_prefix = "abseil-cpp-{commit}".format(commit = ABSL_COMMIT),
-        urls = tf_mirror_urls("https://github.com/abseil/abseil-cpp/archive/{commit}.tar.gz".format(commit = ABSL_COMMIT)),
-    )
diff --git a/third_party/build_option/crosstool_sycl/cc_toolchain_config.bzl.tpl b/third_party/build_option/crosstool_sycl/cc_toolchain_config.bzl.tpl
@@ -39,7 +39,7 @@ load(
     "PREPROCESS_ASSEMBLE_ACTION_NAME",
     "STRIP_ACTION_NAME",
 )
-load("@local_config_sycl//sycl:build_defs.bzl", "if_sycl")
+load("@itex_local_config_sycl//sycl:build_defs.bzl", "if_sycl")
 
 ACTION_NAMES = struct(
     c_compile = C_COMPILE_ACTION_NAME,
@@ -598,6 +598,7 @@ def _impl(ctx):
                         "-fPIC",
                         "-DITEX_USE_MKL=%{TF_NEED_MKL}",
                         "-DITEX_ENABLE_DOUBLE=1",
+                        "-DTENSORFLOW_USE_SYCL=1",
                         "-DEIGEN_USE_DPCPP=1",
                         "-DEIGEN_USE_DPCPP_BUILD=1",
                         "-DEIGEN_USE_DPCPP_USM=1",
diff --git a/third_party/build_option/sycl/build_defs.bzl.tpl b/third_party/build_option/sycl/build_defs.bzl.tpl
@@ -19,7 +19,7 @@ def if_xetla(if_true, if_false = []):
 
     """
     return select({
-        "@local_config_sycl//sycl:using_xetla": if_true,
+        "@itex_local_config_sycl//sycl:using_xetla": if_true,
         "//conditions:default": if_false,
     })     
     
diff --git a/third_party/intel_extension_for_openxla/BUILD b/third_party/intel_extension_for_openxla/BUILD
diff --git a/third_party/intel_extension_for_openxla/intel_extension_for_openxla.BUILD b/third_party/intel_extension_for_openxla/intel_extension_for_openxla.BUILD
diff --git a/third_party/llvm_spir/llvm_spir.patch b/third_party/llvm_spir/llvm_spir.patch
@@ -1,10 +1,10 @@
 diff --git a/lib/SPIRV/SPIRVInternal.h b/lib/SPIRV/SPIRVInternal.h
-index 94a088f0..69403106 100644
+index a828add8..924e13b4 100644
 --- a/lib/SPIRV/SPIRVInternal.h
 +++ b/lib/SPIRV/SPIRVInternal.h
 @@ -179,11 +179,12 @@ typedef SPIRVMap<Op, Op, IntBoolOpMapId> IntBoolOpMap;
    "-v512:512:512-v1024:1024:1024"
- 
+
  enum SPIRAddressSpace {
 -  SPIRAS_Private,
 +  SPIRAS_Generic,
@@ -18,32 +18,3 @@ index 94a088f0..69403106 100644
    SPIRAS_GlobalDevice,
    SPIRAS_GlobalHost,
    SPIRAS_Input,
-diff --git a/lib/SPIRV/SPIRVRegularizeLLVM.cpp b/lib/SPIRV/SPIRVRegularizeLLVM.cpp
-index e9a42beb..9f550bde 100644
---- a/lib/SPIRV/SPIRVRegularizeLLVM.cpp
-+++ b/lib/SPIRV/SPIRVRegularizeLLVM.cpp
-@@ -426,7 +426,6 @@ bool SPIRVRegularizeLLVMBase::regularize() {
-           FI->dropAllReferences();
-           ToErase.push_back(FI);
-         }
--
-         // Remove metadata not supported by SPIRV
-         static const char *MDs[] = {
-             "fpmath",
-diff --git a/lib/SPIRV/SPIRVWriter.cpp b/lib/SPIRV/SPIRVWriter.cpp
-index 90b8c4e2..b049b125 100644
---- a/lib/SPIRV/SPIRVWriter.cpp
-+++ b/lib/SPIRV/SPIRVWriter.cpp
-@@ -367,12 +367,6 @@ SPIRVType *LLVMToSPIRVBase::transType(Type *T) {
-   if (T->isArrayTy()) {
-     // SPIR-V 1.3 s3.32.6: Length is the number of elements in the array.
-     //                     It must be at least 1.
--    if (T->getArrayNumElements() < 1) {
--      std::string Str;
--      llvm::raw_string_ostream OS(Str);
--      OS << *T;
--      SPIRVCK(T->getArrayNumElements() >= 1, InvalidArraySize, OS.str());
--    }
-     Type *ElTy = T->getArrayElementType();
-     SPIRVType *TransType = BM->addArrayType(
-         transType(ElTy),
diff --git a/third_party/onednn/onednn_gpu.BUILD b/third_party/onednn/onednn_gpu.BUILD

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@ cc_library(`
`10`	`10`	`"//itex:core",`
`11`	`11`	`"//itex/core/utils/onednn:onednn_graph_util",`
`12`	`12`	`] + if_using_nextpluggable_device([`
`13`		`- "//itex/core/utils:libintel_xla",`
	`13`	`+ "//itex/core/utils:libintel_extension_for_openxla",`
`14`	`14`	`]),`
`15`	`15`	`alwayslink = True,`
`16`	`16`	`)`