Merge branch 'qiming/main-v0.4.30' into r0.5

Author: Ashiq Imran

README.md

@@ -37,7 +37,7 @@ Verified Hardware Platforms:
 * Ubuntu 22.04, SUSE Linux Enterprise Server(SLES) 15 SP4
   * Intel® Data Center GPU Max Series
 * [Intel® oneAPI Base Toolkit 2024.2](https://www.intel.com/content/www/us/en/developer/articles/release-notes/intel-oneapi-toolkit-release-notes.html)
-* Jax/Jaxlib 0.4.26
+* Jax/Jaxlib 0.4.30
 * Python 3.9-3.12
 * pip 19.0 or later (requires manylinux2014 support)
 

WORKSPACE

@@ -14,10 +14,10 @@ http_archive(
     name = "xla",
     patch_args = ["-p1"],
     patches = ["//third_party:openxla.patch"],
-    sha256 = "fa6e7d17acc362b56c57c43224e6e3eca8569adae864e2fa191cc9d13edf4309",
-    strip_prefix = "xla-4e8e23f16bc925b6f27817de098a8e1e81296bb5",
+    sha256 = "083c7281a629647ab2cc32f054afec74893c33e75328783b8085c818f48235ff",
+    strip_prefix = "xla-79fd5733f99b3c0948d7202bc1bbe1ee3980da5c",
     urls = [
-        "https://github.com/openxla/xla/archive/4e8e23f16bc925b6f27817de098a8e1e81296bb5.tar.gz",
+        "https://github.com/openxla/xla/archive/79fd5733f99b3c0948d7202bc1bbe1ee3980da5c.tar.gz",
     ],
 )
 
@@ -33,6 +33,35 @@ http_archive(
 #    path = "/path/to/xla",
 # )
 
+# Initialize hermetic Python
+load("@xla//third_party/py:python_init_rules.bzl", "python_init_rules")
+
+python_init_rules()
+
+load("@xla//third_party/py:python_init_repositories.bzl", "python_init_repositories")
+
+python_init_repositories(
+    default_python_version = "system",
+    requirements = {
+        "3.9": "@xla//:requirements_lock_3_9.txt",
+        "3.10": "@xla//:requirements_lock_3_10.txt",
+        "3.11": "@xla//:requirements_lock_3_11.txt",
+        "3.12": "@xla//:requirements_lock_3_12.txt",
+    },
+)
+
+load("@xla//third_party/py:python_init_toolchains.bzl", "python_init_toolchains")
+
+python_init_toolchains()
+
+load("@xla//third_party/py:python_init_pip.bzl", "python_init_pip")
+
+python_init_pip()
+
+load("@pypi//:requirements.bzl", "install_deps")
+
+install_deps()
+
 load("@xla//:workspace4.bzl", "xla_workspace4")
 
 xla_workspace4()

example/sdxl/README.md

@@ -19,7 +19,7 @@ please got the [main page](https://github.com/intel/intel-extension-for-openxla/
 Mark `intel-extension-for-openxla` folder as \<WORKSPACE\>, then
 ```bash
 cd <WORKSPACE>/example/sdxl/
-pip install transformers==4.38 diffusers==0.26.3 datasets==2.20.0 msgpack==1.0.7
+pip install transformers==4.47 diffusers==0.31.0 datasets==2.20.0 msgpack==1.1.0
 pip install -r ../../test/requirements.txt
 ```
 

example/t5/install_xpu.sh

@@ -8,7 +8,7 @@ git apply ../patch/t5.patch
 ln -s /usr/local/bin/pip /usr/bin/pip
 pip uninstall tensorflow-metadata numba cudf -y
 pip uninstall tensorflow -y
-pip install tensorflow==2.12.0
+pip install tensorflow==2.18.0
 
 conda install libstdcxx-ng==12.2.0 -c conda-forge -y
 

example/t5/patch/t5.patch

@@ -1,5 +1,5 @@
 diff --git a/setup.py b/setup.py
-index 37238ba..5ee7b8a 100644
+index 37238ba..6a97d34 100644
 --- a/setup.py
 +++ b/setup.py
 @@ -27,8 +27,8 @@ from version import __version__  # pylint: disable=g-import-not-at-top
@@ -8,8 +8,8 @@ index 37238ba..5ee7b8a 100644
 
 -_jax_version = '0.4.11'
 -_jaxlib_version = '0.4.11'
-+_jax_version = '0.4.26'
-+_jaxlib_version = '0.4.26'
++_jax_version = '0.4.30'
++_jaxlib_version = '0.4.30'
 
  setuptools.setup(
      name='t5x',
@@ -19,11 +19,32 @@ index 37238ba..5ee7b8a 100644
          'cached_property',
 -        'clu @ git+https://github.com/google/CommonLoopUtils#egg=clu',
 -        'flax @ git+https://github.com/google/flax#egg=flax',
-+        'clu == 0.0.9',
-+        'flax >= 0.8.2',
++        'clu == 0.0.12',
++        'flax >= 0.8.5',
          'fiddle >= 0.2.5',
          'gin-config',
          f'jax >= {_jax_version}',
+@@ -61,7 +61,7 @@ setuptools.setup(
+         'numpy',
+         'optax @ git+https://github.com/deepmind/optax#egg=optax',
+         'orbax-checkpoint',
+-        'seqio @ git+https://github.com/google/seqio#egg=seqio',
++        'seqio >= 0.0.18',
+         'tensorflow-cpu',
+         'tensorstore >= 0.1.20',
+         # remove this when sentencepiece_model_pb2 is re-generated in the
+diff --git a/t5x/checkpoints.py b/t5x/checkpoints.py
+index c8af7d0..4945b2c 100644
+--- a/t5x/checkpoints.py
++++ b/t5x/checkpoints.py
+@@ -45,7 +45,6 @@ from flax import serialization
+ from flax import traverse_util
+ import jax
+ from jax import monitoring
+-import jax.config
+ from jax.experimental import multihost_utils
+ from jax.experimental.array_serialization import serialization as array_serialization
+ import jax.numpy as jnp
 diff --git a/t5x/config_utils.py b/t5x/config_utils.py
 index abd3f8f..e6e1bd9 100644
 --- a/t5x/config_utils.py

example/t5/quick_start.sh

@@ -65,7 +65,7 @@ echo $MODEL_PATH
 echo "Please make sure ${NUM_GPUS} is the number of visible CUDA devices you have"
 
 # Setting XLA flags
-export XLA_FLAGS="--xla_gpu_simplify_all_fp_conversions --xla_gpu_all_reduce_combine_threshold_bytes=136314880 ${XLA_FLAGS}"
+export XLA_FLAGS="--xla_allow_excess_precision --xla_gpu_all_reduce_combine_threshold_bytes=136314880 ${XLA_FLAGS}"
 
 
 PREFIX=""

test/BRANCH_NAME

@@ -1 +1 @@
-jax-v0.4.26
+jax-v0.4.30

test/requirements.txt

@@ -1,3 +1,3 @@
-jax==0.4.26
-jaxlib==0.4.26
-flax==0.8.2
+jax==0.4.30
+jaxlib==0.4.30
+flax==0.8.5

third_party/onednn/onednn_gpu.BUILD

@@ -35,6 +35,11 @@ _CMAKE_COMMON_LIST = {
     "#cmakedefine DNNL_ENABLE_STACK_CHECKER": "#undef DNNL_ENABLE_STACK_CHECKER",
     "#cmakedefine DNNL_EXPERIMENTAL": "#define DNNL_EXPERIMENTAL",
     "#cmakedefine ONEDNN_BUILD_GRAPH": "#undef ONEDNN_BUILD_GRAPH",
+    "#cmakedefine DNNL_SYCL_GENERIC": "#define DNNL_SYCL_GENERIC",
+    "#cmakedefine DNNL_GPU_VENDOR DNNL_VENDOR_${DNNL_GPU_VENDOR}": "#define DNNL_GPU_VENDOR DNNL_VENDOR_INTEL",
+    "#cmakedefine DNNL_DISABLE_GPU_REF_KERNELS": "#define DNNL_DISABLE_GPU_REF_KERNELS",
+    "#cmakedefine DNNL_EXPERIMENTAL_LOGGING": "#undef DNNL_EXPERIMENTAL_LOGGING",
+    "#cmakedefine01 BUILD_SDPA": "#define BUILD_SDPA 1",
     "#cmakedefine01 BUILD_TRAINING": "#define BUILD_TRAINING 1",
     "#cmakedefine01 BUILD_INFERENCE": "#define BUILD_INFERENCE 0",
     "#cmakedefine01 BUILD_PRIMITIVE_ALL": "#define BUILD_PRIMITIVE_ALL 1",
@@ -103,6 +108,12 @@ gen_onednn_version(
     header_out = "include/oneapi/dnnl/dnnl_version.h",
 )
 
+gen_onednn_version(
+    name = "dnnl_version_hash_h",
+    header_in = "include/oneapi/dnnl/dnnl_version_hash.h.in",
+    header_out = "include/oneapi/dnnl/dnnl_version_hash.h",
+)
+
 filegroup(
     name = "onednn_src",
     srcs = glob(
@@ -122,6 +133,7 @@ filegroup(
         ],
     ) + [
         ":dnnl_config_h",
+	":dnnl_version_hash_h",
         ":header_generator",
         ":kernel_list_generator",
         ":onednn_version_generator",
@@ -156,6 +168,9 @@ cc_library(
         "include/oneapi/dnnl",
         "src",
         "src/common",
+	"src/gpu/intel/jit/gemm/",
+        "src/gpu/intel/jit/gemm/include/",
+        "src/gpu/intel/jit/ngen/",
         "src/intel/ocl",
         "src/sycl",
     ],