Enable armv8 CI (#195)

* [Setup] Skip hatchet pip package for now

This does not exist for Darwin + Arm64.
TODO: Enable this selectively when possible.

* [CPU][driver] Skip non-existent sys paths

* [mac-arm64] Add GH CI support

- look into faster triton install
- enable bf16 tests
- enable openmp

.github/workflows/build-test.yml

@@ -45,14 +45,14 @@ jobs:
           python3 -m pre_commit run --show-diff-on-failure --color=always --all-files --verbose
 
   build-test:
-    name: Build and test
-    runs-on:
-      - glados
-      - intel
-      - x86
+    name: Build and test on ${{ matrix.config.runner }}
+    runs-on: ${{ matrix.config.runs_on }}
     strategy:
       matrix:
         python: ['3.11']
+        config:
+          - {runner: 'Ubuntu Intel x86', runs_on: ['glados', 'intel', 'x86'], target-os: 'ubuntu', arch: 'x86'}
+          - {runner: 'MacOS-latest ARM64', runs_on: ['macos-latest'], target-os: 'macos',  arch: 'arm64'}
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -65,11 +65,16 @@ jobs:
           python-version: ${{ matrix.python }}
 
       - name: Install pip and apt dependencies
+        env:
+          RUNNER_TARGET_OS: ${{ matrix.config.target-os }}
         run: |
+          echo "RUNNER_TARGET_OS: ${RUNNER_TARGET_OS}"
           python3 -m pip install --upgrade pip
           python3 -m pip install wheel cmake==3.24 ninja pytest-xdist lit pybind11
-          sudo apt-get update
-          sudo apt-get install -y zlib1g-dev g++
+          if [[ "${RUNNER_TARGET_OS}" == "ubuntu" ]]; then
+            sudo apt-get update
+            sudo apt-get install -y zlib1g-dev g++
+          fi
           pip install torch==2.1.2
 
       - name: Install Triton
@@ -78,7 +83,49 @@ jobs:
           cd python
           python3 -m pip install --no-build-isolation -vvv '.[tests]'
 
-      - name: Run python unit tests
+      - name: Run python unit tests for MacOS Arm64
+        if: matrix.config.target-os == 'macos'
+        run: |
+          export CC=$(which clang)
+          export TRITON_DISABLE_OPENMP=1 # temporary
+          export TRITON_CPU_BACKEND=1
+
+          # Document some versions/flags
+          echo "xcode-select:"; xcode-select -p
+          echo "CC: ${CC}"
+          clang --version
+          echo "TRITON_DISABLE_OPENMP=${TRITON_DISABLE_OPENMP}"
+          echo "TRITON_CPU_BACKEND=${TRITON_CPU_BACKEND}"
+
+          # Skip bfloat16 tests for now
+          # We are generating bfcvt for bfloat16 tests when converting to fp32.
+          # This is only for Clang15, works OK for Clang16
+          # TODO - fix this using driver flags.
+          python -m pytest -s -n 32 --device cpu \
+            python/test/unit/language/test_core.py -m cpu -k "not bfloat16"
+          python -m pytest -s -n 32 --device cpu \
+            python/test/unit/cpu/test_math.py \
+            python/test/unit/cpu/test_opt.py \
+            python/test/unit/language/test_annotations.py \
+            python/test/unit/language/test_block_pointer.py \
+            python/test/unit/language/test_compile_errors.py \
+            python/test/unit/language/test_conversions.py \
+            python/test/unit/language/test_decorator.py \
+            python/test/unit/language/test_pipeliner.py \
+            python/test/unit/language/test_random.py \
+            python/test/unit/language/test_standard.py \
+            python/test/unit/runtime/test_autotuner.py \
+            python/test/unit/runtime/test_bindings.py \
+            python/test/unit/runtime/test_cache.py \
+            python/test/unit/runtime/test_driver.py \
+            python/test/unit/runtime/test_jit.py \
+            python/test/unit/runtime/test_launch.py \
+            python/test/unit/runtime/test_subproc.py \
+            python/test/unit/test_debug_dump.py \
+            -k "not bfloat16"
+
+      - name: Run python unit tests for Intel
+        if: matrix.config.target-os == 'ubuntu'
         run: |
           python -m pytest -s -n 32 --device cpu python/test/unit/language/test_core.py -m cpu
           python -m pytest -s -n 32 --device cpu \

python/setup.py

@@ -757,7 +757,7 @@ def get_git_commit_hash(length=8):
             "pytest-forked",
             "pytest-xdist",
             "scipy>=1.7.1",
-            "llnl-hatchet",
+            # "llnl-hatchet", # TODO: Re-enable this, not available on macos-arm64
         ],
         "tutorials": [
             "matplotlib",

third_party/cpu/backend/driver.py

@@ -12,6 +12,7 @@
 from triton.backends.driver import DriverBase
 from triton.backends.compiler import GPUTarget
 
+from pathlib import Path
 from triton._C.libtriton import llvm
 
 _dirname = os.getenv("TRITON_SYS_PATH", default="/usr/local")
@@ -22,10 +23,19 @@
     # resources.files() doesn't exist for Python < 3.9
     _triton_C_dir = importlib.resources.path(triton, "_C").__enter__()
 
-include_dirs = [os.path.join(_dirname, "include")]
-library_dirs = [os.path.join(_dirname, "lib"), _triton_C_dir]
+include_dirs = []
+library_dirs = [_triton_C_dir]
 libraries = ["stdc++"]
 
+# Skip non-existent paths
+sys_include_dir = os.path.join(_dirname, "include")
+if os.path.exists(sys_include_dir):
+    include_dirs.append(sys_include_dir)
+
+sys_lib_dir = os.path.join(_dirname, "lib")
+if os.path.exists(sys_lib_dir):
+    library_dirs.append(sys_lib_dir)
+
 
 def compile_module_from_src(src, name):
     key = hashlib.md5(src.encode("utf-8")).hexdigest()