From 4e4e6b8d0a584d1ab26355c8c9b697c85472dbad Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Sat, 21 Dec 2024 02:30:18 -0600 Subject: [PATCH] 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 | 63 +++++++++++++++++++++++++++---- python/setup.py | 2 +- third_party/cpu/backend/driver.py | 14 ++++++- 3 files changed, 68 insertions(+), 11 deletions(-) diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml index a5178e8f34c8..805c6b8dc7b0 100644 --- a/.github/workflows/build-test.yml +++ b/.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 \ diff --git a/python/setup.py b/python/setup.py index 6bf947c10fb6..5f907472770c 100644 --- a/python/setup.py +++ b/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", diff --git a/third_party/cpu/backend/driver.py b/third_party/cpu/backend/driver.py index cadb76e1229a..3308fd23c680 100644 --- a/third_party/cpu/backend/driver.py +++ b/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()