name: CI Test on AMD
on: [pull_request]

env:
  PYTHON_VERSION: '3.12'
  VENV_DIR: tilelang_ci
  PYTORCH_INDEX_URL: https://download.pytorch.org/whl/nightly/rocm6.3/

jobs:
  format-check:
    runs-on: [self-hosted, amd, gpu]

    permissions:
      contents: write 

    steps:
    - name: Checkout repository
      uses: actions/checkout@v4
      with:
        fetch-depth: 0

    - name: Set up Python
      uses: actions/setup-python@v2
      with:
        python-version: ${{ env.PYTHON_VERSION }}

    - name: Ensure venv (local & persistent)
      run: |
        set -e
        REQS_HASH=$(sha256sum requirements-test.txt 2>/dev/null | awk '{print $1}' || echo "no_requirements")
        MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}"

        if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then
          echo "venv exists and hash matches – reuse it"
        else
          echo "venv stale or missing – recreating"
          rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" "$MARKER"
          python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"
          # shellcheck source=/dev/null
          source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
          python -m pip install --upgrade pip --no-user
          [[ -f requirements-test.txt ]] && \
            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-test.txt --no-user
          pip install flash_attn==2.5.8 --no-user --no-build-isolation
          touch "$MARKER"
        fi

    - name: Run format check
      run: |
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        git submodule update --init --recursive
        mkdir -p build
        cd build; cmake .. -DCMAKE_EXPORT_COMPILE_COMMANDS=ON -DUSE_ROCM=ON; cd ..
        if ! output=$(./format.sh 2>&1); then
          echo "------------------------------------"
          echo "message:"
          echo "$output"
          printf '%s\n' "$output" | grep "Please review and stage the changes."
          echo "------------------------------------"
          exit 1
        fi
        rm -rf build

  build-test-amd:
    runs-on: [self-hosted, amd, gpu]
    needs: format-check
    permissions:
      contents: read
    steps:
    - name: Checkout repository
      uses: actions/checkout@v4
      with:
        fetch-depth: 0
        repository: ${{ github.event.pull_request.head.repo.full_name }}
        ref: ${{ github.event.pull_request.head.ref }}

    - name: Set up Python
      uses: actions/setup-python@v2
      with:
        python-version: ${{ env.PYTHON_VERSION }}

    - name: Ensure venv (local & persistent)
      run: |
        echo "Running on AMD GPU"
        set -e
        REQS_HASH=$(sha256sum requirements-rocm.txt | cut -d ' ' -f 1)
        MARKER="${{ runner.tool_cache }}/.venv_marker_${{ env.PYTHON_VERSION }}_${REQS_HASH:0:8}"
        
        echo "Installing requirements"
        if [[ -f "$MARKER" ]] && [[ -f "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate" ]]; then
          echo "venv exists and hash matches – reuse it"
        else
          echo "venv stale or missing – recreating"
          rm -rf "${{ runner.tool_cache }}/${{ env.VENV_DIR }}" "$MARKER"
          python -m venv "${{ runner.tool_cache }}/${{ env.VENV_DIR }}"
          source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
          python -m pip install --upgrade pip --no-user
          if [[ -f requirements-rocm.txt ]]; then
            pip install --pre torch torchvision torchaudio --index-url ${{ env.PYTORCH_INDEX_URL }}
            PIP_NO_BUILD_ISOLATION=1 pip install -r requirements-rocm.txt
          fi

          USE_ROCM=True pip install . --no-user
          touch "$MARKER"
        fi

    - name: Install project (wheel form)
      run: |
        echo "Installing project (wheel form)"
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        USE_ROCM=True pip install . --no-user

    - name: Run tests
      run: |
        echo "Running tests"
        source "${{ runner.tool_cache }}/${{ env.VENV_DIR }}/bin/activate"
        cd testing/python/amd
        unset PYTHONPATH
        python -m pytest -v test_tilelang_test_amd.py