name: Integration Tests AMD

on:
  workflow_call:
    inputs:
      matrix:
        required: true
        type: string

jobs:
  integration-tests-amd:
    runs-on: ${{ matrix.runner }}
    timeout-minutes: 45
    strategy:
      matrix:
        runner: ${{ fromJson(inputs.matrix) }}
    env:
      RUNNER_TYPE: ${{ matrix.runner[1] }}
      TRITON_BUILD_WITH_CCACHE: "true"
      TRITON_BUILD_WITH_CLANG_LLD: "TRUE"
      TRITON_USE_ASSERT_ENABLED_LLVM: "TRUE"
      TRITON_DISABLE_LINE_INFO: 1
      PROTON_SKIP_PC_SAMPLING_TEST: 1
      PYTHON: "python3"
      CCACHE_COMPRESS: "true"
    container:
      image: rocm/pytorch:rocm6.2.2_ubuntu22.04_py3.10_pytorch_2.5.1_asan
      options: --device=/dev/kfd --device=/dev/dri --security-opt seccomp=unconfined --group-add video --user root
    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
          submodules: 'true'
      - name: Compute cache keys
        id: cache-key
        run: |
          llvm_file="cmake/llvm-hash.txt"
          nvidia_file="cmake/nvidia-toolchain-version.json"
          json_file="cmake/json-version.txt"

          # Check if files exist before proceeding
          if [[ ! -f "$llvm_file" || ! -f "$nvidia_file" || ! -f "$json_file" ]]; then
            echo "Error: Required dependency files are missing."
            exit 1
          fi

          # Process the files if they exist
          echo "llvm=$(cat $llvm_file | cut -c 1-8)" >> $GITHUB_OUTPUT
          echo "nvidia=$(sha256sum $nvidia_file | cut -d ' ' -f 1)" >> $GITHUB_OUTPUT
          echo "json=$(cat $json_file)" >> $GITHUB_OUTPUT
          echo "datetime=$(date -u -Iseconds)" >> $GITHUB_OUTPUT
        shell: bash
      - name: Cache build dependencies
        uses: actions/cache@v4
        with:
          # Note that we cannot use environment variables here given there is
          # no shell to interpret them in the paths.
          path: |
            ~/.triton/llvm
            ~/.triton/nvidia
            ~/.triton/json
          key: ${{ runner.os }}-${{ runner.arch }}-llvm-${{ steps.cache-key.outputs.llvm }}-nvidia-${{ steps.cache-key.outputs.nvidia }}-json-${{ steps.cache-key.outputs.json }}
      - # Cache ~/.cache/ccache to speed up compilation.
        #
        # On branch `main` we always start from an empty cache, i.e. we skip the
        # "restore" step.  This is to prevent the caches from accumulating stale
        # files over time.
        name: Restore cache of ccache and Triton compilation artifacts
        id: restore-build-cache
        if: github.ref != 'refs/heads/main'
        uses: actions/cache/restore@v4
        with:
          path: |
            ~/.ccache
          # Restore the most recent cache entry.
          restore-keys: |
            triton-artifacts-${{ runner.os }}-${{ runner.arch }}-${{ env.RUNNER_TYPE }}-llvm-${{ steps.cache-key.outputs.llvm }}-
            triton-artifacts-${{ runner.os }}-${{ runner.arch }}-${{ env.RUNNER_TYPE }}-
          # We expect this cache key never to hit and for us to fall back
          # unconditionally to the restore-key, so it doesn't actually matter
          # what we put here (so long as it doesn't hit an existing key).
          key: triton-artifacts-${{ runner.os }}-${{ runner.arch }}-${{ env.RUNNER_TYPE }}-llvm-${{ steps.cache-key.outputs.llvm }}-${{ steps.cache-key.outputs.datetime }}
      - name: Inspect cache directories
        run: |
          mkdir -p ~/.triton
          du -h -d 1 ~/.triton

          mkdir -p ~/.ccache
          du -h -d 1 ~/.ccache
      - name: Update compiler to clang
        run: |
          export CC=/usr/bin/clang
          export CXX=/usr/bin/clang++
      - name: Install Triton
        id: amd-install-triton
        run: |
          echo "PATH is '$PATH'"
          pip uninstall -y triton pytorch-triton-rocm
          ccache --zero-stats
          make dev-install
      - name: CCache Stats
        run: ccache --print-stats
      - name: Run lit tests
        run: make test-lit
      - name: Run python tests on AMD
        run: |
          INSTRUMENTATION_LIB_DIR="${GITHUB_WORKSPACE}/python/triton/instrumentation"
          if [ ! -d "${INSTRUMENTATION_LIB_DIR}" ]; then
            echo "Could not find '${INSTRUMENTATION_LIB_DIR}'" ; exit -1
          fi
          pytest --capture=tee-sys -rfs python/tutorials/06-fused-attention.py
          pytest --capture=tee-sys -rfs third_party/amd/python/test/test_extract_slice.py
          cd python/test/unit
          pytest --capture=tee-sys -rfs -n 12 language runtime \
                 --ignore=language/test_line_info.py \
                 --ignore=test_debug.py
          # TODO: uncomment
          # pytest --capture=tee-sys -rfs test_debug.py
          TRITON_ALWAYS_COMPILE=1 TRITON_DISABLE_LINE_INFO=0 LLVM_PASS_PLUGIN_PATH=${INSTRUMENTATION_LIB_DIR}/libGPUInstrumentationTestLib.so \
          pytest --capture=tee-sys -rfs -vvv instrumentation/test_gpuhello.py

          # Run test_line_info.py separately with TRITON_DISABLE_LINE_INFO=0
          TRITON_DISABLE_LINE_INFO=0 python3 -m pytest -s -n 8 language/test_line_info.py
      - name: Run asan tests on AMD
        if: false
        run: |
          cd third_party/amd/python/test/
          ulimit -s 1024
          export PATH=$(find ~/.triton/llvm -name llvm-symbolizer  -printf '%h\n'):$PATH
          export LD_LIBRARY_PATH=$(find /opt -name libclang_rt.asan-x86_64.so -printf '%h\n'):$LD_LIBRARY_PATH
          export LD_LIBRARY_PATH=$(find /opt -type d -wholename *lib/llvm/lib/asan):$LD_LIBRARY_PATH
          export LD_LIBRARY_PATH=$(find /usr -name libcaffe2_nvrtc.so -printf '%h\n'):$LD_LIBRARY_PATH
          export CLANG_ASAN_LIB=$(find /opt -name libclang_rt.asan-x86_64.so)
          export HIP_ASAN_LIB=$(find /opt -wholename *lib/asan/libamdhip64.so)
          ASAN_OPTIONS=detect_leaks=0,alloc_dealloc_mismatch=0 \
          LD_PRELOAD=$CLANG_ASAN_LIB:$HIP_ASAN_LIB python3 -m pytest -s test_address_sanitizer.py
      - name: Run regression tests
        run: |
          # Reenable test_functional_regression.py once it's fixed
          cd python/test/regression
          python3 -m pytest -s -n 8 ./test_cast_matmul.py
      - name: Run Proton tests
        if: ${{ matrix.runner[0] != 'nvidia-gb200' }}
        run: make test-proton
      - name: Run C++ unittests
        run: make test-cpp
      - name: Inspect cache directories
        run: |
          mkdir -p ~/.triton
          du -h -d 1 ~/.triton

          mkdir -p ~/.ccache
          du -h -d 1 ~/.ccache
      - # If we're on branch `main`, save the ccache Triton compilation artifacts
        # to the cache so they can be used by other (non-main) CI runs.
        #
        # (It wouldn't be a problem to save the cache on every run, because github
        # evicts cache entries LRU, but maybe this saves a bit of time in CI.)
        name: Save ccache and Triton compilation artifacts to cache
        if: github.ref == 'refs/heads/main'
        uses: actions/cache/save@v4
        with:
          path: |
            ~/.ccache
          key: triton-artifacts-${{ runner.os }}-${{ runner.arch }}-${{ env.RUNNER_TYPE }}-llvm-${{ steps.cache-key.outputs.llvm }}-${{ steps.cache-key.outputs.datetime }}
      - name: Clean up caches
        # Always cleanup the worker, even if builds or tests failed
        if: always()
        run: |
          rm -rf ~/.triton
          rm -rf ~/.ccache
