Add GPU unit tests.

.github/workflows/optim_tests.yml

@@ -34,7 +34,7 @@ jobs:
           export PATH=~/conda/bin:$PATH
           source activate theseus
           pip install torch
-      - name: Build Baspacho GPU
+      - name: Build Baspacho CPU
         run: |
           sudo apt-get install -y libopenblas-pthread-dev
           git clone https://github.com/facebookresearch/baspacho.git

.github/workflows/tests_gpu.yml

@@ -0,0 +1,87 @@
+name: gpu-tests
+on:
+  pull_request:
+    branches:
+    - main
+
+jobs:
+  test_with_gpu:
+    runs-on: 4-core-ubuntu-gpu-t4 # meta configured gpu runner https://www.internalfb.com/intern/opensource/ci/github-actions/runners/
+    strategy:
+      matrix:
+        python-version: [3.10.15]
+    steps:
+      - uses: actions/[email protected]
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/[email protected]
+        with:
+          python-version: ${{ matrix.python-version }}
+      - uses: actions/[email protected]
+      - name: Install suitesparse
+        run: |
+          sudo apt-get update && sudo apt-get install -y libsuitesparse-dev
+      - name: Create Conda env
+        run: |
+          wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh
+          /bin/bash ~/miniconda.sh -b -p ~/conda
+          export PATH=~/conda/bin:$PATH
+          conda create --name theseus python=${{ matrix.python-version }}
+          source activate theseus
+          pip install --progress-bar off --upgrade pip
+          pip install --progress-bar off --upgrade setuptools
+      - name: Install CUDA 11
+        run: |
+          # download and install nvidia drivers, cuda, etc
+          wget --quiet --no-clobber -P ~/nvidia-downloads https://developer.download.nvidia.com/compute/cuda/11.7.1/local_installers/cuda_11.7.1_515.65.01_linux.run
+          time sudo /bin/bash ~/nvidia-downloads/cuda_11.7.1_515.65.01_linux.run --no-drm --silent --driver --toolkit
+          sudo ldconfig /usr/local/cuda/lib64
+          echo "Done installing NVIDIA drivers and CUDA libraries."
+          nvidia-smi
+      - name: Install Torch CUDA 11
+        run: |
+          export PATH=~/conda/bin:$PATH
+          source activate theseus
+          which python && which pip
+          pip install --progress-bar off torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118
+          python -c 'import torch; print("Torch version:", torch.__version__); assert torch.cuda.is_available()'
+      - name: Install recent cmake
+        run: |
+          sudo apt-get update
+          sudo apt-get install ca-certificates gpg wget
+          test -f /usr/share/doc/kitware-archive-keyring/copyright || wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | gpg --dearmor - | sudo tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null
+          echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ jammy main' | sudo tee /etc/apt/sources.list.d/kitware.list >/dev/null
+          sudo apt-get update
+          test -f /usr/share/doc/kitware-archive-keyring/copyright || sudo rm /usr/share/keyrings/kitware-archive-keyring.gpg
+          sudo apt-get install kitware-archive-keyring
+          sudo apt-get install cmake
+      - name: Build Baspacho GPU
+        run: |
+          sudo apt-get install -y libopenblas-pthread-dev
+          git clone https://github.com/facebookresearch/baspacho.git
+          cd baspacho
+          cmake -S . -B build -DCMAKE_BUILD_TYPE=Release -DBLA_STATIC=ON \
+            -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc \
+            -DBUILD_SHARED_LIBS=OFF -DBASPACHO_CUDA_ARCHS='detect' \
+            -DBASPACHO_BUILD_TESTS=OFF -DBASPACHO_BUILD_EXAMPLES=OFF
+          cmake --build build -- -j16
+      - name: Install torchlie and torchkin
+        run:  |
+          export PATH=~/conda/bin:$PATH
+          source activate theseus
+          cd torchlie
+          pip install -e .
+          cd ../torchkin
+          pip install -e .
+      - name: Install theseus
+        run: |
+          export PATH=~/conda/bin:$PATH
+          source activate theseus
+          echo $(pwd)
+          BASPACHO_ROOT_DIR=./baspacho pip install -e ".[dev]"
+      - name: Run Lie groups tests
+        run: |
+          export PATH=~/conda/bin:$PATH
+          source activate theseus
+          pytest -s tests/theseus_tests/test_theseus_layer.py
+          pytest -s tests/theseus_tests/ -m "cudaext"
+          pytest -s tests/theseus_tests/test_pgo_benchmark.py -s