Intel Extension for OpenXLA Containers

.gitignore

@@ -14,4 +14,5 @@ logs/
 models-perf/
 output/
 site
+test-runner-summary-output.json
 venv/

docs/roadmap.md

@@ -10,7 +10,7 @@
 
 - Granite Rapids Support
 - CLS Support
-- Intel Developer Cloud Support
+- Intel Tiber Developer Cloud Support
 - AI Tools 2024.3/2025.0 Support
 
 ## Q4'24

docs/scripts/hook.py

@@ -34,6 +34,7 @@ def create_support_matrix():
     compose_to_csv("pytorch", "serving")
     compose_to_csv("tensorflow", None)
     compose_to_csv("classical-ml", None)
+    compose_to_csv("jax", None)
 
     # get_repo(models)
     compose_to_csv("preset/data-analytics", "data_analytics")

docs/scripts/readmes.py

@@ -17,6 +17,7 @@
 
 readmes = [
     "classical-ml/README.md",
+    "jax/README.md",
     "preset/README.md",
     "python/README.md",
     "pytorch/README.md",

jax/.actions.json

@@ -0,0 +1,5 @@
+{
+    "PACKAGE_OPTION": ["idp", "pip"],
+    "experimental": [true],
+    "runner_label": ["PVC"]
+}

jax/Dockerfile

@@ -0,0 +1,104 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+#
+#
+# This file was assembled from multiple pieces, whose use is documented
+# throughout. Please refer to the TensorFlow dockerfiles documentation
+# for more information.
+
+ARG REGISTRY
+ARG REPO
+ARG GITHUB_RUN_NUMBER
+ARG BASE_IMAGE_NAME
+ARG BASE_IMAGE_TAG
+ARG PACKAGE_OPTION=pip
+ARG PYTHON_VERSION
+ARG PYTHON_BASE=${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER}-${BASE_IMAGE_NAME}-${BASE_IMAGE_TAG}-${PACKAGE_OPTION}-py${PYTHON_VERSION}-base
+ARG TORCHSERVE_BASE=${PYTHON_BASE}
+FROM ${PYTHON_BASE} AS xpu-base
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    apt-utils \
+    build-essential \
+    clinfo \
+    git \
+    gnupg2 \
+    gpg-agent \
+    rsync \
+    unzip && \
+    apt-get clean && \
+    rm -rf  /var/lib/apt/lists/*
+
+RUN wget -qO - https://repositories.intel.com/gpu/intel-graphics.key | \
+    gpg --dearmor --yes --output /usr/share/keyrings/intel-graphics.gpg
+RUN echo "deb [arch=amd64 signed-by=/usr/share/keyrings/intel-graphics.gpg] https://repositories.intel.com/gpu/ubuntu jammy unified" | \
+    tee /etc/apt/sources.list.d/intel-gpu-jammy.list
+
+ARG ICD_VER
+ARG LEVEL_ZERO_GPU_VER
+ARG LEVEL_ZERO_VER
+ARG LEVEL_ZERO_DEV_VER
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    intel-opencl-icd=${ICD_VER} \
+    intel-level-zero-gpu=${LEVEL_ZERO_GPU_VER} \
+    libze1=${LEVEL_ZERO_VER} \
+    libze-dev=${LEVEL_ZERO_DEV_VER} && \
+    rm -rf  /var/lib/apt/lists/*
+
+RUN no_proxy="" NO_PROXY="" wget --progress=dot:giga -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \
+    | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \
+    | tee /etc/apt/sources.list.d/oneAPI.list
+
+ARG DPCPP_VER
+ARG MKL_VER
+ARG CCL_VER
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    intel-oneapi-runtime-dpcpp-cpp=${DPCPP_VER} \
+    intel-oneapi-runtime-mkl=${MKL_VER} \
+    intel-oneapi-runtime-ccl=${CCL_VER} && \
+    rm -rf  /var/lib/apt/lists/*
+
+RUN rm -rf /etc/apt/sources.list.d/intel-gpu-jammy.list /etc/apt/sources.list.d/oneAPI.list
+
+ENV OCL_ICD_VENDORS=/etc/OpenCL/vendors
+
+FROM xpu-base AS jax-base
+
+WORKDIR /
+COPY requirements.txt .
+
+RUN python -m pip install --no-cache-dir -r requirements.txt && \
+    rm -rf requirements.txt
+
+FROM jax-base AS jupyter
+
+WORKDIR /jupyter
+COPY jupyter-requirements.txt .
+
+RUN python -m pip install --no-cache-dir -r jupyter-requirements.txt && \
+    rm -rf jupyter-requirements.txt
+
+RUN mkdir -p /jupyter/ && chmod -R a+rwx /jupyter/
+RUN mkdir /.local && chmod a+rwx /.local
+
+EXPOSE 8888
+
+CMD ["bash", "-c", "source /etc/bash.bashrc && jupyter notebook --notebook-dir=/jupyter --port 8888 --ip 0.0.0.0 --no-browser --allow-root --ServerApp.token= --ServerApp.password= --ServerApp.allow_origin=* --ServerApp.base_url=$NB_PREFIX"]

jax/README.md

@@ -0,0 +1,86 @@
+# Intel® Optimized OpenXLA\*
+
+Transformable numerical computing at scale combined with [Intel® Extension for OpenXLA\*], which includes a PJRT plugin implementation to seamlessly runs [JAX\*] models on Intel GPUs.
+
+## Images
+
+The images below include [JAX\*] and [Intel® Extension for OpenXLA\*].
+
+| Tag(s)                     | [JAX\*]   | [Intel® Extension for OpenXLA\*] | [Flax]   | Dockerfile      |
+| -------------------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-pip-base`, `latest` | [v0.4.32] | [v0.4.0-jax]      | [v0.9.0] | [v0.4.0]        |
+
+The images below additionally include [Jupyter Notebook](https://jupyter.org/) server:
+
+| Tag(s)              | [JAX\*]   | [Intel® Extension for OpenXLA\*] | [Flax]   | Dockerfile      |
+| ------------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-pip-jupyter` | [v0.4.32] | [v0.4.0-jax]      | [v0.9.0] | [v0.4.0]        |
+
+### Run the Jupyter Container
+
+```bash
+docker run -it --rm \
+    -p 8888:8888 \
+    --net=host \
+    -v $PWD/workspace:/workspace \
+    -w /workspace \
+    intel/intel-optimized-xla:0.4.0-pip-jupyter
+```
+
+After running the command above, copy the URL (something like `http://127.0.0.1:$PORT/?token=***`) into your browser to access the notebook server.
+
+## Images with Intel® Distribution for Python*
+
+The images below include [Intel® Distribution for Python*]:
+
+| Tag(s)           | [JAX\*]   | [Intel® Extension for OpenXLA\*] | [Flax]   | Dockerfile      |
+| ---------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-idp-base` | [v0.4.32] | [v0.4.0-jax]      | [v0.9.0] | [v0.4.0]        |
+
+The images below additionally include [Jupyter Notebook](https://jupyter.org/) server:
+
+| Tag(s)              | [JAX\*]   | [Intel® Extension for OpenXLA\*] | [Flax]   | Dockerfile      |
+| ------------------- | --------- | ----------------- | -------- | --------------- |
+| `0.4.0-idp-jupyter` | [v0.4.32] | [v0.4.0-jax]      | [v0.9.0] | [v0.4.0]        |
+
+## Build from Source
+
+To build the images from source, clone the [AI Containers](https://github.com/intel/ai-containers) repository, follow the main `README.md` file to setup your environment, and run the following command:
+
+```bash
+cd jax
+docker compose build jax-base
+docker compose run -it jax-base
+```
+
+You can find the list of services below for each container in the group:
+
+| Service Name | Description                                     |
+| ------------ | ----------------------------------------------- |
+| `jax-base`   | Base image with [Intel® Extension for OpenXLA\*] |
+| `jupyter`    | Adds Jupyter Notebook server                    |
+
+## License
+
+View the [License](https://github.com/intel/ai-containers/blob/main/LICENSE) for the [Intel® Distribution for Python].
+
+The images below also contain other software which may be under other licenses (such as Pytorch*, Jupyter*, Bash, etc. from the base).
+
+It is the image user's responsibility to ensure that any use of The images below comply with any relevant licenses for all software contained within.
+
+\* Other names and brands may be claimed as the property of others.
+
+<!--Below are links used in these document. They are not rendered: -->
+
+[Intel® Distribution for Python*]: https://www.intel.com/content/www/us/en/developer/tools/oneapi/distribution-for-python.html#gs.9bos9m
+[Intel® Extension for OpenXLA\*]: https://github.com/intel/intel-extension-for-openxla
+[JAX\*]: https://github.com/google/jax
+[Flax]: https://github.com/google/flax
+
+[v0.4.32]: https://github.com/google/jax/releases/tag/jax-v0.4.32
+
+[v0.4.0-jax]: https://github.com/intel/intel-extension-for-openxla/releases/tag/0.4.0
+
+[v0.9.0]: https://github.com/google/Flax/releases/tag/v0.9.0
+
+[v0.4.0]: https://github.com/intel/ai-containers/blob/v0.4.0/jax/Dockerfile

jax/docker-compose.yaml

@@ -0,0 +1,87 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+include:
+  - path:
+      - ../python/docker-compose.yaml
+services:
+  jax-base:
+    build:
+      args:
+        http_proxy: ${http_proxy}
+        https_proxy: ${https_proxy}
+        no_proxy: ""
+        BASE_IMAGE_NAME: ${BASE_IMAGE_NAME:-ubuntu}
+        BASE_IMAGE_TAG: ${BASE_IMAGE_TAG:-22.04}
+        CCL_VER: ${CCL_VER:-2021.13.1-31}
+        DPCPP_VER: ${DPCPP_VER:-2024.2.1-1079}
+        GITHUB_RUN_NUMBER: ${GITHUB_RUN_NUMBER:-0}
+        ICD_VER: ${ICD_VER:-24.22.29735.27-914~22.04}
+        LEVEL_ZERO_DEV_VER: ${LEVEL_ZERO_DEV_VER:-1.17.6-914~22.04}
+        LEVEL_ZERO_GPU_VER: ${LEVEL_ZERO_GPU_VER:-1.3.29735.27-914~22.04}
+        LEVEL_ZERO_VER: ${LEVEL_ZERO_VER:-1.17.6-914~22.04}
+        MINIFORGE_VERSION: ${MINIFORGE_VERSION:-Linux-x86_64}
+        MKL_VER: ${MKL_VER:-2024.2.1-103}
+        NO_PROXY: ''
+        PACKAGE_OPTION: ${PACKAGE_OPTION:-pip}
+        PYTHON_VERSION: ${PYTHON_VERSION:-3.10}
+        REGISTRY: ${REGISTRY}
+        REPO: ${REPO}
+      context: .
+      labels:
+        dependency.python: ${PYTHON_VERSION:-3.10}
+        dependency.apt.build-essential: true
+        dependency.apt.clinfo: true
+        dependency.apt.git: true
+        dependency.apt.gnupg2: true
+        dependency.apt.gpg-agent: true
+        dependency.apt.intel-level-zero-gpu: ${LEVEL_ZERO_GPU_VER:-1.3.29735.27-914~22.04}
+        dependency.apt.intel-oneapi-runtime-ccl: ${CCL_VER:-2021.13.1-31}
+        dependency.apt.intel-oneapi-runtime-dpcpp-cpp: ${DPCPP_VER:-2024.2.1-1079}
+        dependency.apt.intel-oneapi-runtime-mkl: ${MKL_VER:-2024.2.1-103}
+        dependency.apt.intel-opencl-icd: ${ICD_VER:-23.43.27642.40-803~22.04}
+        dependency.apt.level-zero: ${LEVEL_ZERO_VER:-1.17.6-914~22.04}
+        dependency.apt.level-zero-dev: ${LEVEL_ZERO_DEV_VER:-1.17.6-914~22.04}
+        dependency.apt.rsync: true
+        dependency.apt.unzip: true
+        dependency.idp.pip: false
+        dependency.python.pip: requirements.txt
+        docs: jax
+        org.opencontainers.base.name: "intel/python:3.10-core"
+        org.opencontainers.image.name: "intel/intel-optimized-xla"
+        org.opencontainers.image.title: "Intel® Optimized XLA Base Image"
+        org.opencontainers.image.version: ${INTEL_XLA_VERSION:-v0.4.0}-${PACKAGE_OPTION:-pip}-base
+      target: jax-base
+    command: >
+      bash -c "python -c 'import jax; print(\"Jax Version:\", jax.__version__)'"
+    depends_on:
+      - ${PACKAGE_OPTION:-pip}
+    image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-base
+    pull_policy: always
+  jupyter:
+    build:
+      labels:
+        dependency.python.pip: jupyter-requirements.txt
+        org.opencontainers.base.name: "intel/intel-optimized-xla:${INTEL_XLA_VERSION:-v0.4.0}-base"
+        org.opencontainers.image.title: "Intel® Optimized XLA Jupyter Base Image"
+        org.opencontainers.image.version: ${INTEL_XLA_VERSION:-v0.4.0}-jupyter
+      target: jupyter
+    command: >
+      bash -c "python -m jupyter --version"
+    environment:
+      http_proxy: ${http_proxy}
+      https_proxy: ${https_proxy}
+    extends: jax-base
+    image: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-jupyter
+    network_mode: host

jax/jupyter-requirements.txt

@@ -0,0 +1,4 @@
+jupyterlab==4.2.4
+jupyterhub==5.1.0
+notebook==7.2.1
+jupyter-server-proxy>=4.1.2

jax/requirements.txt

@@ -0,0 +1,5 @@
+flax==0.8.2
+intel-extension-for-openxla==0.4.0
+jax==0.4.26
+jaxlib==0.4.26
+cython==3.0.11

jax/tests/example.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: skip-file
+
+import jax
+import jax.numpy as jnp
+
+print("jax.local_devices(): ", jax.local_devices())
+
+
+@jax.jit
+def lax_conv():
+    key = jax.random.PRNGKey(0)
+    lhs = jax.random.uniform(key, (2, 1, 9, 9), jnp.float32)
+    rhs = jax.random.uniform(key, (1, 1, 4, 4), jnp.float32)
+    side = jax.random.uniform(key, (1, 1, 1, 1), jnp.float32)
+    out = jax.lax.conv_with_general_padding(
+        lhs, rhs, (1, 1), ((0, 0), (0, 0)), (1, 1), (1, 1)
+    )
+    out = jax.nn.relu(out)
+    out = jnp.multiply(out, side)
+    return out
+
+
+print(lax_conv())

jax/tests/tests.yaml

@@ -0,0 +1,29 @@
+# Copyright (c) 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+---
+jax-import-${PACKAGE_OPTION:-pip}:
+  img: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-base
+  cmd: python -c 'import jax; print("Jax Version:", jax.__version__); print(jax.devices())'
+  device: ["/dev/dri"]
+jax-import-jupyter-${PACKAGE_OPTION:-pip}:
+  img: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-jupyter
+  cmd: sh -c "python -m jupyter --version"
+jax-xpu-example-${PACKAGE_OPTION:-pip}:
+  img: ${REGISTRY}/${REPO}:b-${GITHUB_RUN_NUMBER:-0}-${BASE_IMAGE_NAME:-ubuntu}-${BASE_IMAGE_TAG:-22.04}-${PACKAGE_OPTION:-pip}-py${PYTHON_VERSION:-3.10}-xla-${INTEL_XLA_VERSION:-v0.4.0}-base
+  cmd: python /tests/example.py
+  device: ["/dev/dri"]
+  volumes:
+    - src: $PWD/jax/tests
+      dst: /tests