name: Build and test

on:
  push:
    branches:
      - master
  pull_request:
    branches:
      - master

jobs:
  build-and-test-cpu:
    strategy:
      matrix:
        os: [ubuntu-latest, macos-latest]

    runs-on: ${{ matrix.os }}

    steps:
      - name: Checkout code
        uses: actions/checkout@v3

      - name: Install OpenMP
        run: |
          if [ "${{ runner.os }}" == "Linux" ]; then
            sudo apt-get update && sudo apt-get install -y libomp-dev
          elif [ "${{ runner.os }}" == "macOS" ]; then
            brew install libomp
          fi

      - name: Install dependencies
        run: pip install -r requirements.txt

      - name: Run preprocessing
        run: python prepro_tinyshakespeare.py

      - name: Train model
        run: python train_gpt2.py --device=cpu

      - name: Compile training and testing program
        run: make test_gpt2 train_gpt2

      - name: Execute testing program (With OpenMP)
        run: OMP_NUM_THREADS=8 ./test_gpt2

      - name: Compile training and testing program without OpenMP
        run: NO_OMP=1 make test_gpt2 train_gpt2

      - name: Execute testing program (No OpenMP)
        run: ./test_gpt2

  build-with-cuda-fp32:
    runs-on: ubuntu-latest  # Host OS, Docker will run on top of this
    container:
      image: nvidia/cuda:12.4.1-devel-ubuntu22.04

    steps:
      - name: Checkout code
        uses: actions/checkout@v3

      - name: Build FP32 checkpoint
        run: make train_gpt2fp32cu test_gpt2fp32cu

      - name: Build FP32 precision
        run: PRECISION=FP32 make train_gpt2cu test_gpt2cu profile_gpt2cu

  build-with-cuda-bf16:
    runs-on: ubuntu-latest  # Host OS, Docker will run on top of this
    container:
      image: nvidia/cuda:12.4.1-devel-ubuntu22.04

    steps:
      - name: Checkout code
        uses: actions/checkout@v3

      - name: Build project
        run: PRECISION=BF16 make test_gpt2cu train_gpt2cu profile_gpt2cu

  build-with-cuda-fp16:
    runs-on: ubuntu-latest  # Host OS, Docker will run on top of this
    container:
      image: nvidia/cuda:12.4.1-devel-ubuntu22.04

    steps:
      - name: Checkout code
        uses: actions/checkout@v3

      - name: Build project
        run: PRECISION=FP16 make test_gpt2cu train_gpt2cu profile_gpt2cu