name: Set Matrix

on:
  workflow_call:
    inputs:
      runner-rocm:
        description: 'Override default (linux.rocm.gpu.ecosystem.gfx950.8) ROCm runner label'
        required: false
        type: string
      runner-cuda:
        description: 'Override default (mt-l-x86aavx2-189-704-a10g-8) CUDA runner label'
        required: false
        type: string
      gpu-arch:
        description: "Restrict the matrix to a single GPU arch ('cuda' or 'rocm'); 'all' (default) keeps both."
        required: false
        type: string
        default: all
      is-experimental:
        description: 'Check if the workflow is experimental'
        required: false
        type: boolean
        default: false
    outputs:
      matrix:
        description: dynamically set matrix, narrowed to the gpu-arch input (both CUDA and ROCm when 'all')
        value: ${{ jobs.set.outputs.matrix }}
      docker-hash:
        description: Git tree hash of .ci/docker, used to build the OSDC image tag
        value: ${{ jobs.set.outputs.docker-hash }}

jobs:
  set:
    if: github.event_name != 'pull_request' || github.event.pull_request.draft == false
    runs-on: ubuntu-latest
    outputs:
      matrix: ${{ steps.set.outputs.matrix }}
      docker-hash: ${{ steps.set.outputs.docker-hash }}
    env:
      # Event flags evaluated by github actions before the step runs:
      IS_MAIN_PUSH: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
      IS_SCHEDULE:  ${{ github.event_name == 'schedule' }}

      # Default CUDA & ROCm runners
      DEFAULT_ROCM_RUNNER: linux.rocm.gpu.ecosystem.mi350.8
      DEFAULT_CUDA_RUNNER: mt-l-x86aavx2-189-704-a10g-8

      # Input CUDA & ROCm runners
      ROCM_RUNNER_INPUT: ${{ inputs.runner-rocm }}
      CUDA_RUNNER_INPUT: ${{ inputs.runner-cuda }}

    steps:
      - uses: actions/checkout@v7

      - id: set
        run: |
          set -euo pipefail

          # .ci/docker tree hash, used to build the OSDC image tag.
          DOCKER_HASH=$(git rev-parse HEAD:.ci/docker)
          echo "docker-hash=${DOCKER_HASH}" >> "$GITHUB_OUTPUT"

          # Select runner if provided, else use default value
          ROCM_RUNNER="${ROCM_RUNNER_INPUT:-$DEFAULT_ROCM_RUNNER}"
          CUDA_RUNNER="${CUDA_RUNNER_INPUT:-$DEFAULT_CUDA_RUNNER}"

          # Define CUDA matrix
          CUDA_MATRIX="$(cat <<EOF
          {
            "name": "cuda",
            "runner": "${CUDA_RUNNER}",
            "gpu-arch-type": "cuda",
            "gpu-arch-version": "13.2",
            "docker-image": "torchtitan-ubuntu-22.04-clang12",
            "index-url": "https://download.pytorch.org/whl/nightly/cu132",
            "torch-version": ""
          }
          EOF
          )"

          if [[ "${{ inputs.is-experimental }}" == "true" ]]; then
            ROCM_MATRIX=""
            FULL_MATRIX="$CUDA_MATRIX"
          else
            ROCM_MATRIX="$(cat <<EOF
          {
            "name": "rocm",
            "runner": "${ROCM_RUNNER}",
            "gpu-arch-type": "rocm",
            "gpu-arch-version": "10.0",
            "docker-image": "torchtitan-rocm-ubuntu-22.04-clang12",
            "index-url": "https://download.pytorch.org/whl/nightly/rocm10.0",
            "torch-version": ""
          }
          EOF
            )"
            FULL_MATRIX="$CUDA_MATRIX,$ROCM_MATRIX"
          fi

          # Use default value as 'false' for unset environment variables
          IS_MAIN_PUSH="${IS_MAIN_PUSH:-false}"
          IS_SCHEDULE="${IS_SCHEDULE:-false}"

          # Decide which matrix entries to include based on event type.
          # Runs CUDA and ROCm for push to main, cron schedules, or when
          # gpu-arch is explicitly set to rocm (e.g. ciflow/rocm/* tag runs).
          # Runs CUDA only as default, including pull requests.
          GPU_ARCH_INPUT="${{ inputs.gpu-arch }}"
          if [[ "$IS_MAIN_PUSH" == "true" || "$IS_SCHEDULE" == "true" || "${GPU_ARCH_INPUT}" == "rocm" ]]; then
            cat > matrix.json <<JSON
          {"include": [$FULL_MATRIX]}
          JSON
          else
            cat > matrix.json <<JSON
          {"include": [$CUDA_MATRIX]}
          JSON
          fi

          # Optionally narrow the matrix to a single GPU arch. Each consuming
          # workflow runs one arch on its own reusable workflow (the H100 CUDA
          # workflow on linux_job_v3, the ROCm workflow on linux_job_v2 -- ROCm
          # must stay on v2 because v3 does not authenticate ROCm runners to the
          # private ECR). 'all' (the default) keeps both for the workflows that
          # run a single combined linux_job_v2 job.
          case "${GPU_ARCH_INPUT:-all}" in
            all) ;;
            cuda|rocm)
              jq -c --arg arch "$GPU_ARCH_INPUT" '{include: [.include[] | select(.["gpu-arch-type"] == $arch)]}' matrix.json > matrix.filtered.json
              mv matrix.filtered.json matrix.json
              ;;
            *)
              echo "::error::Invalid gpu-arch '${GPU_ARCH_INPUT}'. Expected one of: all, cuda, rocm." >&2
              exit 1
              ;;
          esac

          # Export matrix to job outputs
          {
            echo 'matrix<<EOF'
            cat matrix.json
            echo 'EOF'
          } >> $GITHUB_OUTPUT
