diff --git a/.github/workflows/building_xpu.yml b/.github/workflows/building_xpu.yml new file mode 100644 index 00000000..f272c3c6 --- /dev/null +++ b/.github/workflows/building_xpu.yml @@ -0,0 +1,122 @@ +name: Build XPU Wheels + +on: [workflow_call, workflow_dispatch] + +permissions: + contents: read + +jobs: + build_sdist: + name: Build source distribution and no binary wheel + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + submodules: recursive + - name: Strip unsupported tags in README + run: | + sed -i '//,//d' README.md + - name: Build sdist + run: BUILD_NO_CUDA=1 pipx run build --sdist + - name: Build wheel with no binaries + run: BUILD_NO_CUDA=1 python setup.py bdist_wheel --dist-dir=dist + - uses: actions/upload-artifact@v4 + with: + name: pypi_packages + path: dist/*.tar.gz + + build_wheels: + runs-on: ${{ matrix.os }} + environment: production + + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04, windows-2022] + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] + torch-version: ['2.10.0'] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Free up disk space + if: ${{ runner.os == 'Linux' }} + run: | + echo "Disk space before cleanup:" + df -h + sudo rm -rf /usr/share/dotnet + echo "Disk space after cleanup:" + df -h + shell: bash + + - name: Install PyTorch ${{ matrix.torch-version }}+xpu + run: | + pip install torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/xpu + python -c "import torch; print('PyTorch:', torch.__version__)" + python -c "import torch; print('XPU Available:', torch.xpu.is_available() if hasattr(torch, 'xpu') else False)" + shell: bash + + - name: Install Intel oneAPI C++ Essentials + run: | + ONEAPI_VERSION=$(pip show intel-sycl-rt 2>/dev/null | grep ^Version | awk '{print $2}') + if [ -z "${ONEAPI_VERSION}" ]; then + echo "Error: intel-sycl-rt not found. Ensure PyTorch XPU was installed successfully." >&2 + exit 1 + fi + echo "Detected intel-sycl-rt version: ${ONEAPI_VERSION}" + bash .github/workflows/xpu/${RUNNER_OS}.sh ${ONEAPI_VERSION} + shell: bash + + - name: Set version + run: | + VERSION=`sed -n 's/^__version__ = "\(.*\)"/\1/p' gsplat/version.py` + TORCH_VERSION=`echo "pt${{ matrix.torch-version }}" | sed "s/..$//" | sed "s/\.//g"` + echo "New version name: $VERSION+${TORCH_VERSION}xpu" + sed -i "s/$VERSION/$VERSION+${TORCH_VERSION}xpu/" gsplat/version.py + shell: bash + + - name: Upgrade pip + run: | + pip install --upgrade setuptools wheel + pip install ninja pybind11 + shell: bash + + - name: Build wheel (Windows) + if: ${{ runner.os == 'Windows' }} + shell: cmd + run: | + set MAX_JOBS=%NUMBER_OF_PROCESSORS% + set BUILD_SYCL=1 + call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" + python setup.py bdist_wheel --dist-dir=dist + + - name: Build wheel (Linux) + if: ${{ runner.os == 'Linux' }} + shell: bash + run: | + export MAX_JOBS=$(nproc) BUILD_SYCL=1 + source /opt/intel/oneapi/setvars.sh + python setup.py bdist_wheel --dist-dir=dist + + - name: Test wheel + run: | + cd dist + ls -lah + pip install *.whl + python -c "import gsplat; print('gsplat:', gsplat.__version__)" + cd .. + shell: bash + + - uses: actions/upload-artifact@v4 + with: + # Include unique matrix values to avoid name collisions. + name: xpu_wheels_python${{ matrix.python-version }}-${{ matrix.os }}-${{ matrix.torch-version }} + path: dist/*.whl diff --git a/.github/workflows/core_tests.yml b/.github/workflows/core_tests.yml index a4db904f..97fde27f 100644 --- a/.github/workflows/core_tests.yml +++ b/.github/workflows/core_tests.yml @@ -18,15 +18,15 @@ jobs: with: submodules: 'recursive' - - name: Set up Python 3.8.12 + - name: Set up Python 3.9 uses: actions/setup-python@v5 with: - python-version: "3.8.12" + python-version: "3.9" - name: Install dependencies run: | pip install black[jupyter]==22.3.0 pytest - pip install torch==2.0.0 --index-url https://download.pytorch.org/whl/cpu - BUILD_NO_CUDA=1 pip install . + pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu + BUILD_NO_CUDA=1 pip install --no-build-isolation . - name: Run Black Format Check run: black . gsplat/ tests/ examples/ profiling/ --check - name: Run Tests. diff --git a/.github/workflows/doc.yml b/.github/workflows/doc.yml index c70c1060..31e97142 100644 --- a/.github/workflows/doc.yml +++ b/.github/workflows/doc.yml @@ -26,8 +26,8 @@ jobs: - name: Install dependencies run: | pip install -r docs/requirements.txt - pip install torch==2.0.0 --index-url https://download.pytorch.org/whl/cpu - BUILD_NO_CUDA=1 pip install . + pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cpu + BUILD_NO_CUDA=1 pip install --no-build-isolation . # Get version. - name: Get version + subdirectory diff --git a/.github/workflows/generate_simple_index_pages.yml b/.github/workflows/generate_simple_index_pages.yml index 17b3397d..81563916 100644 --- a/.github/workflows/generate_simple_index_pages.yml +++ b/.github/workflows/generate_simple_index_pages.yml @@ -1,5 +1,7 @@ # This workflows will upload a Python Package using twine when a release is created # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries +permissions: + contents: write name: Update wheels index pages diff --git a/.github/workflows/publish_xpu.yml b/.github/workflows/publish_xpu.yml new file mode 100644 index 00000000..986cf8a3 --- /dev/null +++ b/.github/workflows/publish_xpu.yml @@ -0,0 +1,88 @@ +# Build and Release XPU Wheels + +name: Build and Release XPU Wheels + +on: + release: + types: [created] + workflow_dispatch: + +permissions: + contents: write + +jobs: + # Build the XPU wheels using the reusable building workflow + build_xpu_wheels: + name: Call reusable XPU building workflow + uses: ./.github/workflows/building_xpu.yml + + create_release_and_upload_packages: + name: Upload XPU Wheels to GitHub Release + needs: [build_xpu_wheels] + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + submodules: recursive + + - name: Download packages + id: download_artifacts + uses: actions/download-artifact@v4 + with: + # The unique artifact names from building_xpu.yml all start with + # "xpu_wheels_python${{ matrix.python-version }}" so this pattern + # will match them all and merge them into the 'dist' directory. + pattern: xpu_wheels_python${{ matrix.python-version }}* + path: dist + merge-multiple: true + + - name: Upload packages to latest GitHub Release + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + echo "Fetching latest release info..." + release_info=$(curl -s -H "Authorization: token $GITHUB_TOKEN" \ + https://api.github.com/repos/isl-org/gsplat/releases/latest) + + # Extract the "upload_url" field and strip the {?name,label} part + upload_url=$(echo "$release_info" | grep '"upload_url":' | cut -d '"' -f 4 | sed 's/{.*//') + echo "Upload URL: $upload_url" + + for file in ./dist/*.*; do + echo "Uploading $file..." + filename=$(basename "$file") + encoded_filename=$(echo "$filename" | sed 's/+/%2B/g') + curl -X POST \ + -H "Authorization: token $GITHUB_TOKEN" \ + -H "Content-Type: application/octet-stream" \ + --data-binary @"$file" \ + "$upload_url?name=$encoded_filename" + done + echo "Upload complete." + + generate_simple_index_pages: + name: Generate Simple Index Pages + needs: [create_release_and_upload_packages] + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Generate Simple Index Pages + run: python .github/workflows/generate_simple_index_pages.py --outdir ./whl + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + - name: Deploy to GitHub Pages + uses: peaceiris/actions-gh-pages@v3 + with: + github_token: ${{ secrets.GITHUB_TOKEN }} + publish_dir: ./whl + destination_dir: whl + keep_files: false + cname: docs.gsplat.studio diff --git a/.github/workflows/xpu/Linux.sh b/.github/workflows/xpu/Linux.sh new file mode 100644 index 00000000..eb6a3da2 --- /dev/null +++ b/.github/workflows/xpu/Linux.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +# Install Intel oneAPI C++ Essentials for SYCL builds. +# Usage: Linux.sh +# Example: Linux.sh 2025.3.1 + +VERSION=${1:?'Usage: Linux.sh (e.g. Linux.sh 2025.3.1)'} + +# The apt package uses X.Y version format (e.g. 2025.3), while the pip package +# (intel-sycl-rt) may use X.Y.Z format (e.g. 2025.3.1). Truncate to X.Y. +APT_VERSION=$(echo "${VERSION}" | grep -oP '^\d+\.\d+') +if [ -z "${APT_VERSION}" ]; then + echo "Error: VERSION '${VERSION}' does not match expected X.Y or X.Y.Z format." >&2 + exit 1 +fi + +wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \ + | gpg --dearmor \ + | sudo tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null + +echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" \ + | sudo tee /etc/apt/sources.list.d/oneAPI.list + +sudo apt-get -qq update +sudo apt-get install -y intel-cpp-essentials-${APT_VERSION} +sudo apt clean diff --git a/.github/workflows/xpu/Windows.sh b/.github/workflows/xpu/Windows.sh new file mode 100644 index 00000000..d9e3e90d --- /dev/null +++ b/.github/workflows/xpu/Windows.sh @@ -0,0 +1,42 @@ +#!/bin/bash + +# Install Intel oneAPI C++ Essentials for SYCL builds on Windows. +# Usage: Windows.sh +# Example: Windows.sh 2025.3.1 + +set -euo pipefail + +VERSION=${1:?'Usage: Windows.sh (e.g. Windows.sh 2025.3.1)'} + +# Lookup table of Intel C++ Essentials *online* installer URLs per oneAPI version. +# URLs are obtained from (select Windows / Online Installer): +# https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html?packages=cpp-essentials&cpp-essentials-os=windows&cpp-essentials-win=online +# To add a new version, append: ["X.Y.Z"]="https://registrationcenter-download.intel.com/..." +declare -A INSTALLER_URLS=( + ["2025.1.0"]="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/1e635719-29c5-4775-8252-268d2f87d529/intel-cpp-essentials-2025.1.0.570.exe" + ["2025.1"]="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/1e635719-29c5-4775-8252-268d2f87d529/intel-cpp-essentials-2025.1.0.570.exe" + ["2025.2.0"]="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5b271b39-0773-49a3-b78d-c73ec42d1621/intel-cpp-essentials-2025.2.0.533.exe" + ["2025.2"]="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/5b271b39-0773-49a3-b78d-c73ec42d1621/intel-cpp-essentials-2025.2.0.533.exe" + ["2025.3.1"]="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/c61634af-e4dd-4a14-8341-0b35a9ebc22e/intel-cpp-essentials-2025.3.1.25.exe" + ["2025.3"]="https://registrationcenter-download.intel.com/akdlm/IRC_NAS/c61634af-e4dd-4a14-8341-0b35a9ebc22e/intel-cpp-essentials-2025.3.1.25.exe" +) + +INSTALLER_URL="${INSTALLER_URLS[${VERSION}]:-}" +if [[ -z "${INSTALLER_URL}" ]]; then + echo "Error: No installer URL found for oneAPI version '${VERSION}'." >&2 + echo "Add it to the INSTALLER_URLS table in $(basename "${BASH_SOURCE[0]}")." >&2 + echo "Download page: https://www.intel.com/content/www/us/en/developer/tools/oneapi/base-toolkit-download.html?packages=cpp-essentials&cpp-essentials-os=windows&cpp-essentials-win=online" >&2 + exit 1 +fi + +# Install only compiler + necessary libraries. +ONEAPI_WINDOWS_COMPONENTS="${ONEAPI_WINDOWS_COMPONENTS:-intel.oneapi.win.cpp-dpcpp-common}" +INSTALLER_FILE="w_cpp-essentials_p_${VERSION}.exe" +echo "Downloading Intel C++ Essentials online installer from: ${INSTALLER_URL}" +curl -fL "${INSTALLER_URL}" --output "${INSTALLER_FILE}" + +#https://www.intel.com/content/www/us/en/docs/oneapi/installation-guide-windows/2025-2/base-command-line-options.html#BASE-COMMAND-LINE-OPTIONS +echo "Installing components: ${ONEAPI_WINDOWS_COMPONENTS}" +PowerShell -NoProfile -Command "\$p = Start-Process -FilePath '${INSTALLER_FILE}' -ArgumentList '--a -s --action install --eula accept --components ${ONEAPI_WINDOWS_COMPONENTS}' -Wait -PassThru -NoNewWindow; exit \$p.ExitCode" + +rm "${INSTALLER_FILE}" \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in index 16e9cc94..88ee995a 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,8 @@ recursive-include gsplat/cuda/csrc * +recursive-exclude gsplat/cuda/csrc/third_party/glm/doc * +recursive-exclude gsplat/cuda/csrc/third_party/glm/test * recursive-include gsplat/cuda/include * -include gsplat/cuda/ext.cpp \ No newline at end of file +include gsplat/cuda/ext.cpp +recursive-include gsplat/sycl/src * +recursive-include gsplat/sycl/include * +include gsplat/sycl/ext.cpp \ No newline at end of file diff --git a/README.md b/README.md index a73d0952..a7dbca18 100644 --- a/README.md +++ b/README.md @@ -41,6 +41,8 @@ pip install ninja numpy jaxtyping rich pip install gsplat --index-url https://docs.gsplat.studio/whl/pt20cu118 ``` +For Intel XPU (integrated and discrete GPU) support, see [Intel XPU](docs/Intel_XPU.md). + To build gsplat from source on Windows, please check [this instruction](docs/INSTALL_WIN.md). ## Evaluation diff --git a/docs/Intel_XPU.md b/docs/Intel_XPU.md new file mode 100644 index 00000000..2a6b818f --- /dev/null +++ b/docs/Intel_XPU.md @@ -0,0 +1,194 @@ +# GSPLAT on Intel GPUs + +`gsplat` supports creation and rendering on Intel GPUs through the SYCL kernel backend. This provides support for both integrated (Alder Lake Arc and onward) and discrete GPUs (Arc Alchemist and newer, such as the A770 and B580). + +## Supported Features: + +- [x] 3DGS fused training +- [x] 3DGS packed representation +- [x] Distributed training (PyTorch 2.8+) +- [x] MCMC strategy (relocation kernel) +- [x] 2DGS fused training +- [ ] 2DGS packed representation +- [ ] 3DGUT kernels (+FTheta cameras) +- [ ] Fused Bilateral grid kernels (from https://github.com/harry7557558/fused-bilagrid) +- [ ] 3DGS Compression (requires PLAS) +- [ ] `rasterize_to_indices_{2,3}dgs` and `rasterize_to_pixels_from_world_3dgs` kernels (only used internally for testing) + +The kernels are optimized and use mixed precision (some data is represented as half), so the results differ slightly from the CUDA kernel results. + +## Installing (Linux or Windows): + +- **PyTorch XPU:** Install the PyTorch XPU version. + + ```bash + pip install torch torchvision --index-url https://download.pytorch.org/whl/xpu + ``` + + Next install gsplat-xpu directly from here (for PyTorch 2.10+xpu). Otherwise, you can build and install from source. + + ```bash + pip install gsplat --find-links https://isl-org.github.io/gsplat/whl/gsplat + ``` + +- **Intel oneAPI Toolkit:** Ensure you have the [Intel oneAPI Toolkit installed](https://www.intel.com/content/www/us/en/developer/articles/guide/installation-guide-for-oneapi-toolkits.html). This provides the necessary compilers and libraries for SYCL development. + + **Note:** The OneAPI toolkit version must match the version used to build PyTorch XPU. Check the PyTorch XPU OneAPI version with: + + ```bash + pip show intel-cmplr-lib-ur # dependency of torch-xpu + # ... + # Version: 2025.3.1 + # ... + ``` + +- Configure your build environment: + + In Linux: + + ```bash + source /opt/intel/oneapi/setvars.sh + ``` + + Or in Windows, setup your Visual Studio build environment and then OneAPI build environment. For example: + + ```ps1 + cmd /k "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" + powershell + $env:DISTUTILS_USE_SDK=1 + ``` + +- Finally, build and install the project's Python extension. + + ```bash + pip install --extra-index-url=https://download.pytorch.org/whl/xpu . + ``` + + Alternately, you can build a wheel for distribution with: + + ```bash + PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/xpu python -m build --no-isolation --wheel . + ``` + +## Evaluation + +We evaluate gsplat-xpu on the Mip-NeRF 360 dataset and measure PSNR, SSIM, LPIPS and the number of Gaussians used. We also measure the memory used and the run time on an Intel Arc B580 dGPU and an Intel Arc B390 iGPU. To run the evaluation yourself, download the MIPS-NeRF 360 dataset and install other requirements: + + ```bash + cd examples + pip install --extra-index-url=https://download.pytorch.org/whl/xpu -r requirements_xpu.txt + # download mipnerf_360 benchmark data + python datasets/download_dataset.py + ``` + +The last command will also build and install the `fused-ssim` package. Before running benchmarks, you can add `--max-steps 7000` to each `simple_trainer.py` command in `benchmarks/basic{,_2dgs}.sh`, if you have limited memory, or want to run the training faster. Run the benchmarks with: + + ```bash + # run batch evaluation + bash benchmarks/basic.sh + bash benchmarks/basic_2dgs.sh + ``` + +### Arc B580 dGPU + +#### 3DGS Reproduced metrics + +| PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 24.01 | 29.66 | 27.26 | 26.59 | 28.65 | 28.70 | 26.03 | +| 30k steps | [^1] | 31.89 | 29.14 | [^1] | 30.90 | 31.06 | [^1] | + + +| SSIM | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 0.6808 | 0.9262 | 0.8865 | 0.8370 | 0.9047 | 0.8945| 0.7378| +| 30k steps | [^1] | 0.9446 | 0.9158 | [^1] | 0.9318 | 0.9239| [^1] | + + +| LPIPS | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 0.2997 | 0.1462 | 0.1929 | 0.1195 | 0.1220 | 0.2136| 0.2339| +| 30k steps | [^1] | 0.1179 | 0.1414 | [^1] | 0.08607 | 0.1520| [^1] | + +| Num GSs | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 3.95 M | 1.19 M | 1.06 M | 4.20 M | 1.77 M | 1.14 M| 4.04 M| +| 30k steps | [^1] | 1.28 M | 1.27 M | [^1] | 1.90 M | 1.63 M| [^1] | + +#### 3DGS Training time and memory + +| Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------------------|---------|--------|---------|--------|---------|--------|--------| +| 7k steps Mem (GB) | 5.846 | 2.009 | 1.730 | 6.177 | 2.737 | 1.861 | 5.921 | +| 30k steps Mem (GB) | [^1] | 2.043 | 1.986 | [^1] | 2.923 | 2.463 | [^1] | +| 7k steps time (s) | 588.5 | 534.4 | 625.3 | 793.6 | 919.1 | 645.1 | 519.3 | +| 30k steps time (s) | [^1] | 2612 | 3520 | [^1] | 5027 | 3317 | [^1] | + +#### 2DGS Reproduced metrics + +| PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 23.59 | 29.73 | 27.25 | 26.31 | 29.02 | 29.56 | 25.69 | +| 30k steps | 25.33 | 32.13 | 28.90 | 27.39 | 31.33 | 31.43 | 26.72 | + + +| SSIM | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 0.6578 | 0.9277 | 0.8819 | 0.8222 | 0.9021 | 0.9026| 0.7225| +| 30k steps | 0.7570 | 0.9453 | 0.9097 | 0.8567 | 0.9283 | 0.9249| 0.7743| + + +| LPIPS | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 0.3091 | 0.1441 | 0.1935 | 0.1289 | 0.1231 | 0.1988| 0.2403| +| 30k steps | 0.1745 | 0.1173 | 0.1503 | 0.08466| 0.09162 | 0.1555| 0.1565| + +| Num GSs | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | 2.52 M | 0.911 M| 0.695 M | 2.18 M | 0.856 M | 0.839 M| 2.69 M| +| 30k steps | 3.67 M | 0.929 M| 0.731 M | 2.39 M | 0.870 M | 1.03 M| 3.30 M| + +#### 2DGS Training time and memory + +| Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------------------|---------|--------|---------|--------|---------|--------|--------| +| 7k steps Mem (GB) | 4.621 | 2.129 | 1.832 | 4.004 | 2.063 | 2.057 | 4.766 | +| 30k steps Mem (GB) | 6.491 | 2.129 | 1.854 | 4.278 | 2.063 | 2.224 | 5.802 | +| 7k steps time (s) | 560.0 | 758.2 | 666.3 | 609.3 | 732.8 | 643.8 | 545.6 | +| 30k steps time (s) | 3483 | 3308 | 2941 | 3101 | 3196 | 2936 | 3117 | + +[^1]: Out of memory. + +### Arc B390 iGPU + +#### 3DGS Reproduced metrics + +| 7k steps | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| PSNR | 24.02 | 29.72 | 27.40 | 26.61 | 29.24 | 29.56 | 25.31 | +| SSIM | 0.6513 | 0.9252 | 0.8914 | 0.8369 | 0.9173 | 0.9039| 0.6955| +| LPIPS | 0.3558 | 0.1525 | 0.1891 | 0.1198 | 0.1102 | 0.2037| 0.2941| +| Num GSs | 3.28 M | 0.99 M | 0.74 M | 4.17 M | 1.07 M | 0.80 M| 4.01 M| + +#### 3DGS Training time and memory + +| Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|--------------------|---------|--------|---------|--------|---------|--------|--------| +| 7k steps Mem (GB) | 5.016 | 1.642 | 1.264 | 6.142 | 1.678 | 1.366 | 5.932 | +| 7k steps time (s) | 1346.5 | 1124.1 | 1231.0 | 1958.6 | 1496.2 | 983.1 | 1215.6 | + +#### 2DGS Reproduced metrics + +| 7k steps | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| PSNR | 23.92 | 29.88 | 27.38 | 25.95 | 29.37 | 29.98 | 25.13 | +| SSIM | 0.6418 | 0.9301 | 0.8897 | 0.7992 | 0.9128 | 0.9086| 0.6855| +| LPIPS | 0.3443 | 0.1446 | 0.1843 | 0.1520 | 0.1123 | 0.1919| 0.2902| +| Num GSs | 2.13 M | 0.79 M | 0.56 M | 1.66 M | 0.72 M | 0.62 M| 2.40 M| + +#### 2DGS Training time and memory + +| Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------------------|---------|--------|---------|--------|---------|--------|--------| +| 7k steps Mem (GB) | 3.026 | 1.9185 | 1.5837 | 3.026 | 1.8087 | 1.6828 | 4.263 | +| 7k steps time (s) | 1502.5 | 1868.3 | 2121.1 | 1502.5 | 1816.6 | 1591.3 | 1409.4 | \ No newline at end of file diff --git a/docs/source/apis/utils.rst b/docs/source/apis/utils.rst index cc839dab..4a701244 100644 --- a/docs/source/apis/utils.rst +++ b/docs/source/apis/utils.rst @@ -27,8 +27,6 @@ Below are the basic functions that supports the rasterization. .. autofunction:: rasterize_to_indices_in_range -.. autofunction:: accumulate - .. autofunction:: rasterization_inria_wrapper 2DGS @@ -41,6 +39,4 @@ Below are the basic functions that supports the rasterization. .. autofunction:: rasterize_to_indices_in_range_2dgs -.. autofunction:: accumulate_2dgs - .. autofunction:: rasterization_2dgs_inria_wrapper \ No newline at end of file diff --git a/examples/benchmarks/basic.sh b/examples/benchmarks/basic.sh index 6b043c56..c4e90d6f 100644 --- a/examples/benchmarks/basic.sh +++ b/examples/benchmarks/basic.sh @@ -16,7 +16,7 @@ do # train without eval CUDA_VISIBLE_DEVICES=0 python simple_trainer.py default --eval_steps -1 --disable_viewer --data_factor $DATA_FACTOR \ --render_traj_path $RENDER_TRAJ_PATH \ - --data_dir data/360_v2/$SCENE/ \ + --data_dir $SCENE_DIR/$SCENE/ \ --result_dir $RESULT_DIR/$SCENE/ # run eval and render @@ -24,7 +24,7 @@ do do CUDA_VISIBLE_DEVICES=0 python simple_trainer.py default --disable_viewer --data_factor $DATA_FACTOR \ --render_traj_path $RENDER_TRAJ_PATH \ - --data_dir data/360_v2/$SCENE/ \ + --data_dir $SCENE_DIR/$SCENE/ \ --result_dir $RESULT_DIR/$SCENE/ \ --ckpt $CKPT done @@ -44,7 +44,7 @@ do echo "=== Train Stats ===" - for STATS in $RESULT_DIR/$SCENE/stats/train*_rank0.json; + for STATS in $RESULT_DIR/$SCENE/stats/train*.json; do echo $STATS cat $STATS; diff --git a/examples/benchmarks/basic_2dgs.sh b/examples/benchmarks/basic_2dgs.sh index 04d3d8fd..b8881838 100755 --- a/examples/benchmarks/basic_2dgs.sh +++ b/examples/benchmarks/basic_2dgs.sh @@ -15,7 +15,7 @@ do # train without eval CUDA_VISIBLE_DEVICES=0 python simple_trainer_2dgs.py --eval_steps -1 --disable_viewer --data_factor $DATA_FACTOR \ --model_type 2dgs \ - --data_dir data/360_v2/$SCENE/ \ + --data_dir $SCENE_DIR/$SCENE/ \ --result_dir $RESULT_DIR/$SCENE/ # run eval and render @@ -23,7 +23,7 @@ do do CUDA_VISIBLE_DEVICES=0 python simple_trainer_2dgs.py --disable_viewer --data_factor $DATA_FACTOR \ --model_type 2dgs \ - --data_dir data/360_v2/$SCENE/ \ + --data_dir $SCENE_DIR/$SCENE/ \ --result_dir $RESULT_DIR/$SCENE/ \ --ckpt $CKPT done @@ -43,7 +43,7 @@ do echo "=== Train Stats ===" - for STATS in $RESULT_DIR/$SCENE/stats/train*_rank0.json; + for STATS in $RESULT_DIR/$SCENE/stats/train*.json; do echo $STATS cat $STATS; diff --git a/examples/benchmarks/basic_4gpus.sh b/examples/benchmarks/basic_4gpus.sh index 283c583b..81afd3ec 100644 --- a/examples/benchmarks/basic_4gpus.sh +++ b/examples/benchmarks/basic_4gpus.sh @@ -17,7 +17,7 @@ do # "--packed" reduces the data transfer between GPUs, which leads to faster training. CUDA_VISIBLE_DEVICES=0,1,2,3 python simple_trainer.py default --eval_steps 30000 --disable_viewer --data_factor $DATA_FACTOR \ --steps_scaler 0.25 --packed \ - --data_dir data/360_v2/$SCENE/ \ + --data_dir $SCENE_DIR/$SCENE/ \ --result_dir $RESULT_DIR/$SCENE/ done diff --git a/examples/image_fitting.py b/examples/image_fitting.py index 434b7869..2344ca28 100644 --- a/examples/image_fitting.py +++ b/examples/image_fitting.py @@ -10,7 +10,7 @@ from PIL import Image from torch import Tensor, optim -from gsplat import rasterization, rasterization_2dgs +from gsplat import torch_acc, rasterization, rasterization_2dgs class SimpleTrainer: @@ -21,7 +21,7 @@ def __init__( gt_image: Tensor, num_points: int = 2000, ): - self.device = torch.device("cuda:0") + self.device = torch_acc._get_device(0) self.gt_image = gt_image.to(device=self.device) self.num_points = num_points @@ -117,13 +117,13 @@ def train( packed=False, )[0] out_img = renders[0] - torch.cuda.synchronize() + torch_acc.synchronize() times[0] += time.time() - start loss = mse_loss(out_img, self.gt_image) optimizer.zero_grad() start = time.time() loss.backward() - torch.cuda.synchronize() + torch_acc.synchronize() times[1] += time.time() - start optimizer.step() print(f"Iteration {iter + 1}/{iterations}, Loss: {loss.item()}") diff --git a/examples/requirements.txt b/examples/requirements.txt index ea0a940e..a55535fb 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -19,6 +19,6 @@ tensorboard tensorly pyyaml matplotlib -git+https://github.com/rahul-goel/fused-ssim@328dc9836f513d00c4b5bc38fe30478b4435cbb5 +git+https://github.com/rahul-goel/fused-ssim@b42e988db507702aa1198920ec7b36a5aa3f72b2 git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe splines diff --git a/examples/requirements_xpu.txt b/examples/requirements_xpu.txt new file mode 100644 index 00000000..73b3951e --- /dev/null +++ b/examples/requirements_xpu.txt @@ -0,0 +1,24 @@ +# assume torch is already installed + +# pycolmap for data parsing +git+https://github.com/rmbrualla/pycolmap@cc7ea4b7301720ac29287dbe450952511b32125e +# (optional) nerfacc for torch version rasterization +# git+https://github.com/nerfstudio-project/nerfacc + +viser +git+https://github.com/nerfstudio-project/nerfview@4538024fe0d15fd1a0e4d760f3695fc44ca72787 +imageio[ffmpeg] +numpy<2.0.0 +scikit-learn +tqdm +torchmetrics[image] +opencv-python +tyro>=0.8.8 +Pillow +tensorboard +tensorly +pyyaml +matplotlib +git+https://github.com/rahul-goel/fused-ssim@b42e988db507702aa1198920ec7b36a5aa3f72b2 +#git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe +splines diff --git a/examples/simple_trainer.py b/examples/simple_trainer.py index 3f73f0e5..0f46d2b2 100644 --- a/examples/simple_trainer.py +++ b/examples/simple_trainer.py @@ -21,6 +21,7 @@ generate_interpolated_path, generate_spiral_path, ) + from fused_ssim import fused_ssim from torch import Tensor from torch.nn.parallel import DistributedDataParallel as DDP @@ -30,7 +31,7 @@ from typing_extensions import Literal, assert_never from utils import AppearanceOptModule, CameraOptModule, knn, rgb_to_sh, set_random_seed -from gsplat import export_splats +from gsplat import export_splats, torch_acc, BACKEND from gsplat.compression import PngCompression from gsplat.distributed import cli from gsplat.optimizers import SelectiveAdam @@ -79,13 +80,13 @@ class Config: # Number of training steps max_steps: int = 30_000 # Steps to evaluate the model - eval_steps: List[int] = field(default_factory=lambda: [7_000, 30_000]) + eval_steps: List[int] = field(default_factory=lambda: [2_000, 7_000, 30_000]) # Steps to save the model - save_steps: List[int] = field(default_factory=lambda: [7_000, 30_000]) + save_steps: List[int] = field(default_factory=lambda: [2_000, 7_000, 30_000]) # Whether to save ply file (storage size can be large) save_ply: bool = False # Steps to save the model as ply - ply_steps: List[int] = field(default_factory=lambda: [7_000, 30_000]) + ply_steps: List[int] = field(default_factory=lambda: [2_000, 7_000, 30_000]) # Whether to disable video generation during training and evaluation disable_video: bool = False @@ -227,7 +228,7 @@ def create_splats_with_optimizers( visible_adam: bool = False, batch_size: int = 1, feature_dim: Optional[int] = None, - device: str = "cuda", + device: str = torch_acc._device(0).type, world_rank: int = 0, world_size: int = 1, ) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]: @@ -294,7 +295,7 @@ def create_splats_with_optimizers( eps=1e-15 / math.sqrt(BS), # TODO: check betas logic when BS is larger than 10 betas[0] will be zero. betas=(1 - BS * (1 - 0.9), 1 - BS * (1 - 0.999)), - fused=True, + fused=(None if BACKEND == "sycl" else True), ) for name, _, lr in params } @@ -313,7 +314,7 @@ def __init__( self.world_rank = world_rank self.local_rank = local_rank self.world_size = world_size - self.device = f"cuda:{local_rank}" + self.device = str(torch_acc._device(local_rank)) # Where to dump results. os.makedirs(cfg.result_dir, exist_ok=True) @@ -736,7 +737,7 @@ def train(self): # ) if world_rank == 0 and cfg.tb_every > 0 and step % cfg.tb_every == 0: - mem = torch.cuda.max_memory_allocated() / 1024**3 + mem = torch_acc.max_memory_allocated() / 1024**3 self.writer.add_scalar("train/loss", loss.item(), step) self.writer.add_scalar("train/l1loss", l1loss.item(), step) self.writer.add_scalar("train/ssimloss", ssimloss.item(), step) @@ -754,7 +755,7 @@ def train(self): # save checkpoint before updating the model if step in [i - 1 for i in cfg.save_steps] or step == max_steps - 1: - mem = torch.cuda.max_memory_allocated() / 1024**3 + mem = torch_acc.max_memory_allocated() / 1024**3 stats = { "mem": mem, "ellipse_time": time.time() - global_tic, @@ -924,7 +925,7 @@ def eval(self, step: int, stage: str = "val"): masks = data["mask"].to(device) if "mask" in data else None height, width = pixels.shape[1:3] - torch.cuda.synchronize() + torch_acc.synchronize() tic = time.time() colors, _, _ = self.rasterize_splats( camtoworlds=camtoworlds, @@ -936,7 +937,7 @@ def eval(self, step: int, stage: str = "val"): far_plane=cfg.far_plane, masks=masks, ) # [1, H, W, 3] - torch.cuda.synchronize() + torch_acc.synchronize() ellipse_time += max(time.time() - tic, 1e-10) colors = torch.clamp(colors, 0.0, 1.0) @@ -1178,6 +1179,37 @@ def main(local_rank: int, world_rank, world_size: int, cfg: Config): step = ckpts[0]["step"] runner.eval(step=step) runner.render_traj(step=step) + if cfg.save_ply: + if runner.cfg.app_opt: + # eval at origin to bake the appeareance into the colors + rgb = runner.app_module( + features=runner.splats["features"], + embed_ids=None, + dirs=torch.zeros_like(runner.splats["means"][None, :, :]), + sh_degree=runner.cfg.sh_degree, + ) + rgb = rgb + runner.splats["colors"] + rgb = torch.sigmoid(rgb).squeeze(0).unsqueeze(1) + sh0 = rgb_to_sh(rgb) + shN = torch.empty([sh0.shape[0], 0, 3], device=sh0.device) + else: + sh0 = runner.splats["sh0"] + shN = runner.splats["shN"] + + means = runner.splats["means"] + scales = runner.splats["scales"] + quats = runner.splats["quats"] + opacities = runner.splats["opacities"] + export_splats( + means=means, + scales=scales, + quats=quats, + opacities=opacities, + sh0=sh0, + shN=shN, + format="ply", + save_to=f"{cfg.result_dir}/point_cloud_{step}.ply", + ) if cfg.compression is not None: runner.run_compression(step=step) else: diff --git a/examples/simple_trainer_2dgs.py b/examples/simple_trainer_2dgs.py index fcca5993..2afda112 100644 --- a/examples/simple_trainer_2dgs.py +++ b/examples/simple_trainer_2dgs.py @@ -29,6 +29,7 @@ rgb_to_sh, set_random_seed, ) +from gsplat import torch_acc from gsplat_viewer_2dgs import GsplatViewer, GsplatRenderTabState from gsplat.rendering import rasterization_2dgs, rasterization_2dgs_inria_wrapper from gsplat.strategy import DefaultStrategy @@ -194,7 +195,7 @@ def create_splats_with_optimizers( sparse_grad: bool = False, batch_size: int = 1, feature_dim: Optional[int] = None, - device: str = "cuda", + device: str = torch_acc._device(0).type, ) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]: if init_type == "sfm": points = torch.from_numpy(parser.points).float() @@ -257,7 +258,7 @@ def __init__(self, cfg: Config) -> None: set_random_seed(42) self.cfg = cfg - self.device = "cuda" + self.device = torch_acc._device(0).type # Where to dump results. os.makedirs(cfg.result_dir, exist_ok=True) @@ -650,7 +651,7 @@ def train(self): pbar.set_description(desc) if cfg.tb_every > 0 and step % cfg.tb_every == 0: - mem = torch.cuda.max_memory_allocated() / 1024**3 + mem = torch_acc.max_memory_allocated() / 1024**3 self.writer.add_scalar("train/loss", loss.item(), step) self.writer.add_scalar("train/l1loss", l1loss.item(), step) self.writer.add_scalar("train/ssimloss", ssimloss.item(), step) @@ -712,7 +713,7 @@ def train(self): # save checkpoint if step in [i - 1 for i in cfg.save_steps] or step == max_steps - 1: - mem = torch.cuda.max_memory_allocated() / 1024**3 + mem = torch_acc.max_memory_allocated() / 1024**3 stats = { "mem": mem, "ellipse_time": time.time() - global_tic, @@ -765,7 +766,7 @@ def eval(self, step: int): pixels = data["image"].to(device) / 255.0 height, width = pixels.shape[1:3] - torch.cuda.synchronize() + torch_acc.synchronize() tic = time.time() ( colors, @@ -787,7 +788,7 @@ def eval(self, step: int): ) # [1, H, W, 3] colors = torch.clamp(colors, 0.0, 1.0) colors = colors[..., :3] # Take RGB channels - torch.cuda.synchronize() + torch_acc.synchronize() ellipse_time += max(time.time() - tic, 1e-10) # write images diff --git a/examples/simple_viewer.py b/examples/simple_viewer.py index b46746c6..7186861a 100644 --- a/examples/simple_viewer.py +++ b/examples/simple_viewer.py @@ -20,7 +20,7 @@ def main(local_rank: int, world_rank, world_size: int, args): torch.manual_seed(42) - device = torch.device("cuda", local_rank) + device = torch.device(local_rank) if args.ckpt is None: ( diff --git a/examples/utils.py b/examples/utils.py index 80f8e35f..b3cae3fd 100644 --- a/examples/utils.py +++ b/examples/utils.py @@ -88,7 +88,7 @@ def forward( Returns: colors: (C, N, 3) """ - from gsplat.cuda._torch_impl import _eval_sh_bases_fast + from gsplat._torch_impl import _eval_sh_bases_fast C, N = dirs.shape[:2] # Camera embeddings diff --git a/formatter.sh b/formatter.sh index e2747f78..84b9e7ff 100644 --- a/formatter.sh +++ b/formatter.sh @@ -4,5 +4,9 @@ find gsplat/cuda/include \ -type f \( -iname "*.cpp" -o -iname "*.cuh" -o -iname "*.cu" -o -iname "*.h" \) \ -exec clang-format -i {} \; +find gsplat/sycl \ + -type f \( -iname "*.cpp" -o -iname "*.hpp" \) \ + -exec clang-format -i {} \; + # install via: pip install black==22.3.0 black . gsplat/ tests/ examples/ profiling/ \ No newline at end of file diff --git a/gsplat/__init__.py b/gsplat/__init__.py index 27c7b0a9..febfbe3d 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -1,9 +1,20 @@ -import warnings +import os +import sys +import torch -from .compression import PngCompression -from .cuda._torch_impl import accumulate -from .cuda._torch_impl_2dgs import accumulate_2dgs -from .cuda._wrapper import ( +BACKEND: str = "" +torch_acc = torch.cpu + +if torch.cuda.is_available(): + BACKEND = "cuda" + torch_acc = torch.cuda + print("gsplat: Using CUDA backend.", file=sys.stderr) +elif hasattr(torch, "xpu") and torch.xpu.is_available(): + BACKEND = "sycl" + torch_acc = torch.xpu + print("gsplat: Using SYCL XPU backend.", file=sys.stderr) + +from ._wrapper import ( RollingShutterType, fully_fused_projection, fully_fused_projection_2dgs, @@ -20,6 +31,7 @@ spherical_harmonics, world_to_cam, ) +from .compression import PngCompression from .exporter import export_splats from .optimizers import SelectiveAdam from .rendering import ( @@ -31,7 +43,9 @@ from .strategy import DefaultStrategy, MCMCStrategy, Strategy from .version import __version__ -all = [ +__all__ = [ + "BACKEND", + "torch_acc", "PngCompression", "DefaultStrategy", "MCMCStrategy", @@ -47,16 +61,16 @@ "quat_scale_to_covar_preci", "rasterize_to_pixels", "world_to_cam", - "accumulate", "rasterize_to_indices_in_range", "fully_fused_projection_2dgs", "rasterize_to_pixels_2dgs", "rasterize_to_indices_in_range_2dgs", - "accumulate_2dgs", "rasterization_2dgs_inria_wrapper", "RollingShutterType", "fully_fused_projection_with_ut", "rasterize_to_pixels_eval3d", "export_splats", "__version__", + "SelectiveAdam", + # Note: accumulate and accumulate_2dgs are not typically part of the public API ] diff --git a/gsplat/_helper.py b/gsplat/_helper.py index 86bca052..c02818b8 100644 --- a/gsplat/_helper.py +++ b/gsplat/_helper.py @@ -8,7 +8,7 @@ def load_test_data( data_path: Optional[str] = None, - device="cuda", + device=None, scene_crop: Tuple[float, float, float, float, float, float] = (-2, -2, -2, 2, 2, 2), scene_grid: int = 1, ): @@ -19,6 +19,12 @@ def load_test_data( data_path = os.path.join(os.path.dirname(__file__), "../assets/test_garden.npz") data = np.load(data_path) height, width = data["height"].item(), data["width"].item() + if device is None: + device = ( + torch.accelerator.current_accelerator() + if torch.accelerator.is_available() + else torch.device("cpu") + ) viewmats = torch.from_numpy(data["viewmats"]).float().to(device) Ks = torch.from_numpy(data["Ks"]).float().to(device) means = torch.from_numpy(data["means3d"]).float().to(device) diff --git a/gsplat/cuda/_torch_impl.py b/gsplat/_torch_impl.py similarity index 98% rename from gsplat/cuda/_torch_impl.py rename to gsplat/_torch_impl.py index 29888d3d..f06ca9df 100644 --- a/gsplat/cuda/_torch_impl.py +++ b/gsplat/_torch_impl.py @@ -49,7 +49,7 @@ def _quat_scale_to_covar_preci( compute_preci: bool = True, triu: bool = False, ) -> Tuple[Optional[Tensor], Optional[Tensor]]: - """PyTorch implementation of `gsplat.cuda._wrapper.quat_scale_to_covar_preci()`.""" + """PyTorch implementation of `gsplat._wrapper.quat_scale_to_covar_preci()`.""" batch_dims = quats.shape[:-1] assert quats.shape == batch_dims + (4,), quats.shape assert scales.shape == batch_dims + (3,), scales.shape @@ -296,7 +296,7 @@ def _fully_fused_projection( calc_compensations: bool = False, camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Optional[Tensor]]: - """PyTorch implementation of `gsplat.cuda._wrapper.fully_fused_projection()` + """PyTorch implementation of `gsplat._wrapper.fully_fused_projection()` .. note:: @@ -384,7 +384,7 @@ def _isect_tiles( tile_height: int, sort: bool = True, ) -> Tuple[Tensor, Tensor, Tensor]: - """Pytorch implementation of `gsplat.cuda._wrapper.isect_tiles()`. + """Pytorch implementation of `gsplat._wrapper.isect_tiles()`. .. note:: @@ -477,7 +477,7 @@ def kernel(image_id, gauss_id): def _isect_offset_encode( isect_ids: Tensor, I: int, tile_width: int, tile_height: int ) -> Tensor: - """Pytorch implementation of `gsplat.cuda._wrapper.isect_offset_encode()`. + """Pytorch implementation of `gsplat._wrapper.isect_offset_encode()`. .. note:: @@ -617,7 +617,7 @@ def _rasterize_to_pixels( backgrounds: Optional[Tensor] = None, # [..., channels] batch_per_iter: int = 100, ): - """Pytorch implementation of `gsplat.cuda._wrapper.rasterize_to_pixels()`. + """Pytorch implementation of `gsplat._wrapper.rasterize_to_pixels()`. This function rasterizes 2D Gaussians to pixels in a Pytorch-friendly way. It iteratively accumulates the renderings within each batch of Gaussians. The @@ -806,7 +806,7 @@ def _spherical_harmonics( dirs: torch.Tensor, # [..., 3] coeffs: torch.Tensor, # [..., K, 3] ): - """Pytorch implementation of `gsplat.cuda._wrapper.spherical_harmonics()`.""" + """Pytorch implementation of `gsplat._wrapper.spherical_harmonics()`.""" assert (degrees_to_use + 1) ** 2 <= coeffs.shape[-2], coeffs.shape batch_dims = dirs.shape[:-1] assert dirs.shape == batch_dims + (3,), dirs.shape diff --git a/gsplat/cuda/_torch_impl_2dgs.py b/gsplat/_torch_impl_2dgs.py similarity index 98% rename from gsplat/cuda/_torch_impl_2dgs.py rename to gsplat/_torch_impl_2dgs.py index 96ae8695..7f3a0ab8 100644 --- a/gsplat/cuda/_torch_impl_2dgs.py +++ b/gsplat/_torch_impl_2dgs.py @@ -4,7 +4,7 @@ import torch from torch import Tensor -from gsplat.cuda._torch_impl import _quat_scale_to_matrix +from ._torch_impl import _quat_scale_to_matrix def _fully_fused_projection_2dgs( @@ -19,7 +19,7 @@ def _fully_fused_projection_2dgs( far_plane: float = 1e10, eps: float = 0, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """PyTorch implementation of `gsplat.cuda._wrapper.fully_fused_projection_2dgs()` + """PyTorch implementation of `gsplat._wrapper.fully_fused_projection_2dgs()` .. note:: @@ -209,7 +209,7 @@ def _rasterize_to_pixels_2dgs( backgrounds: Optional[Tensor] = None, # [..., channels] batch_per_iter: int = 100, ): - """Pytorch implementation of `gsplat.cuda._wrapper.rasterize_to_pixels_2dgs()`. + """Pytorch implementation of `gsplat._wrapper.rasterize_to_pixels_2dgs()`. This function rasterizes 2D Gaussians to pixels in a Pytorch-friendly way. It iteratively accumulates the renderings within each batch of Gaussians. The diff --git a/gsplat/cuda/_wrapper.py b/gsplat/_wrapper.py similarity index 96% rename from gsplat/cuda/_wrapper.py rename to gsplat/_wrapper.py index f45e2762..72ae3f55 100644 --- a/gsplat/cuda/_wrapper.py +++ b/gsplat/_wrapper.py @@ -8,20 +8,26 @@ from torch import Tensor from typing_extensions import Literal +from . import torch_acc, BACKEND -def _make_lazy_cuda_func(name: str) -> Callable: - def call_cuda(*args, **kwargs): - # pylint: disable=import-outside-toplevel - from ._backend import _C + +def _make_lazy_device_func(name: str) -> Callable: + def call_device(*args, **kwargs): + if BACKEND == "cuda": + from .cuda._backend import _C + elif BACKEND == "sycl": + from .sycl._backend import _C return getattr(_C, name)(*args, **kwargs) - return call_cuda + return call_device -def _make_lazy_cuda_obj(name: str) -> Any: - # pylint: disable=import-outside-toplevel - from ._backend import _C +def _make_lazy_device_obj(name: str) -> Any: + if BACKEND == "cuda": + from .cuda._backend import _C + elif BACKEND == "sycl": + from .sycl._backend import _C obj = _C for name_split in name.split("."): @@ -37,7 +43,7 @@ class RollingShutterType(Enum): GLOBAL = 4 def to_cpp(self) -> Any: - return _make_lazy_cuda_obj(f"ShutterType.{self.name}") + return _make_lazy_device_obj(f"ShutterType.{self.name}") @dataclass @@ -53,7 +59,7 @@ class UnscentedTransformParameters: require_all_sigma_points_valid: bool = True def to_cpp(self) -> Any: - p = _make_lazy_cuda_obj("UnscentedTransformParameters")() + p = _make_lazy_device_obj("UnscentedTransformParameters")() p.alpha = self.alpha p.beta = self.beta p.kappa = self.kappa @@ -68,7 +74,7 @@ class FThetaPolynomialType(Enum): ANGLE_TO_PIXELDIST = 1 def to_cpp(self) -> Any: - return _make_lazy_cuda_obj(f"FThetaPolynomialType.{self.name}") + return _make_lazy_device_obj(f"FThetaPolynomialType.{self.name}") @dataclass @@ -80,7 +86,7 @@ class FThetaCameraDistortionParameters: linear_cde: Tuple[float, float, float] # [3] def to_cpp(self) -> Any: - p = _make_lazy_cuda_obj("FThetaCameraDistortionParameters")() + p = _make_lazy_device_obj("FThetaCameraDistortionParameters")() p.reference_poly = self.reference_poly.to_cpp() p.pixeldist_to_angle_poly = self.pixeldist_to_angle_poly p.angle_to_pixeldist_poly = self.angle_to_pixeldist_poly @@ -90,7 +96,7 @@ def to_cpp(self) -> Any: @classmethod def to_cpp_default(cls) -> Any: - p = _make_lazy_cuda_obj("FThetaCameraDistortionParameters")() + p = _make_lazy_device_obj("FThetaCameraDistortionParameters")() return p @@ -115,7 +121,7 @@ def world_to_cam( from ._torch_impl import _world_to_cam warnings.warn( - "world_to_cam() is removed from the CUDA backend as it's relatively easy to " + "world_to_cam() is removed from the device backend as it's relatively easy to " "implement in PyTorch. Currently use the PyTorch implementation instead. " "This function will be completely removed in a future release.", DeprecationWarning, @@ -143,7 +149,7 @@ def adam( b2: float, eps: float, ) -> None: - _make_lazy_cuda_func("adam")( + _make_lazy_device_func("adam")( param, param_grad, exp_avg, exp_avg_sq, valid, lr, b1, b2, eps ) @@ -317,14 +323,14 @@ def fully_fused_projection( an indicator, in which zero radii means the corresponding elements are invalid in the output tensors and will be ignored in the next rasterization process. If `packed=True`, the output tensors will be packed into a flattened tensor, in which all elements are valid. - In this case, a ``batch_ids` tensor and `camera_ids` tensor will be returned to indicate the + In this case, a `batch_ids` tensor and `camera_ids` tensor will be returned to indicate the batch, camera and gaussian indices of the packed flattened tensor, which is essentially following the COO sparse tensor format. .. note:: This functions supports projecting Gaussians with either covariances or {quaternions, scales}, - which will be converted to covariances internally in a fused CUDA kernel. Either `covars` or + which will be converted to covariances internally in a fused device kernel. Either `covars` or {`quats`, `scales`} should be provided. Args: @@ -502,7 +508,7 @@ def isect_tiles( assert radii.shape == image_dims + (N, 2), radii.shape assert depths.shape == image_dims + (N,), depths.shape - tiles_per_gauss, isect_ids, flatten_ids = _make_lazy_cuda_func("intersect_tile")( + tiles_per_gauss, isect_ids, flatten_ids = _make_lazy_device_func("intersect_tile")( means2d.contiguous(), radii.contiguous(), depths.contiguous(), @@ -536,7 +542,7 @@ def isect_offset_encode( Returns: Offsets. [I, tile_height, tile_width] """ - return _make_lazy_cuda_func("intersect_offset")( + return _make_lazy_device_func("intersect_offset")( isect_ids.contiguous(), n_images, tile_width, tile_height ) @@ -915,7 +921,7 @@ def rasterize_to_indices_in_range( tile_width * tile_size >= image_width ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - out_gauss_ids, out_indices = _make_lazy_cuda_func("rasterize_to_indices_3dgs")( + out_gauss_ids, out_indices = _make_lazy_device_func("rasterize_to_indices_3dgs")( range_start, range_end, transmittances.contiguous(), @@ -945,7 +951,7 @@ def forward( compute_preci: bool = True, triu: bool = False, ) -> Tuple[Tensor, Tensor]: - covars, precis = _make_lazy_cuda_func("quat_scale_to_covar_preci_fwd")( + covars, precis = _make_lazy_device_func("quat_scale_to_covar_preci_fwd")( quats, scales, compute_covar, compute_preci, triu ) ctx.save_for_backward(quats, scales) @@ -964,7 +970,7 @@ def backward(ctx, v_covars: Tensor, v_precis: Tensor): v_covars = v_covars.to_dense() if compute_preci and v_precis.is_sparse: v_precis = v_precis.to_dense() - v_quats, v_scales = _make_lazy_cuda_func("quat_scale_to_covar_preci_bwd")( + v_quats, v_scales = _make_lazy_device_func("quat_scale_to_covar_preci_bwd")( quats, scales, triu, @@ -991,11 +997,11 @@ def forward( camera_model != "ftheta" ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - camera_model_type = _make_lazy_cuda_obj( + camera_model_type = _make_lazy_device_obj( f"CameraModelType.{camera_model.upper()}" ) - means2d, covars2d = _make_lazy_cuda_func("projection_ewa_simple_fwd")( + means2d, covars2d = _make_lazy_device_func("projection_ewa_simple_fwd")( means, covars, Ks, @@ -1015,7 +1021,7 @@ def backward(ctx, v_means2d: Tensor, v_covars2d: Tensor): width = ctx.width height = ctx.height camera_model_type = ctx.camera_model_type - v_means, v_covars = _make_lazy_cuda_func("projection_ewa_simple_bwd")( + v_means, v_covars = _make_lazy_device_func("projection_ewa_simple_bwd")( means, covars, Ks, @@ -1054,12 +1060,12 @@ def forward( camera_model != "ftheta" ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - camera_model_type = _make_lazy_cuda_obj( + camera_model_type = _make_lazy_device_obj( f"CameraModelType.{camera_model.upper()}" ) # "covars" and {"quats", "scales"} are mutually exclusive - radii, means2d, depths, conics, compensations = _make_lazy_cuda_func( + radii, means2d, depths, conics, compensations = _make_lazy_device_func( "projection_ewa_3dgs_fused_fwd" )( means, @@ -1109,7 +1115,7 @@ def backward(ctx, v_radii, v_means2d, v_depths, v_conics, v_compensations): camera_model_type = ctx.camera_model_type if v_compensations is not None: v_compensations = v_compensations.contiguous() - v_means, v_covars, v_quats, v_scales, v_viewmats = _make_lazy_cuda_func( + v_means, v_covars, v_quats, v_scales, v_viewmats = _make_lazy_device_func( "projection_ewa_3dgs_fused_bwd" )( means, @@ -1215,9 +1221,9 @@ def fully_fused_projection_with_ut( if viewmats_rs is not None: assert viewmats_rs.shape == batch_dims + (C, 4, 4), viewmats_rs.shape - camera_model_type = _make_lazy_cuda_obj(f"CameraModelType.{camera_model.upper()}") + camera_model_type = _make_lazy_device_obj(f"CameraModelType.{camera_model.upper()}") - radii, means2d, depths, conics, compensations = _make_lazy_cuda_func( + radii, means2d, depths, conics, compensations = _make_lazy_device_func( "projection_ut_3dgs_fused" )( means.contiguous(), @@ -1240,9 +1246,11 @@ def fully_fused_projection_with_ut( radial_coeffs.contiguous() if radial_coeffs is not None else None, tangential_coeffs.contiguous() if tangential_coeffs is not None else None, thin_prism_coeffs.contiguous() if thin_prism_coeffs is not None else None, - ftheta_coeffs.to_cpp() - if ftheta_coeffs is not None - else FThetaCameraDistortionParameters.to_cpp_default(), + ( + ftheta_coeffs.to_cpp() + if ftheta_coeffs is not None + else FThetaCameraDistortionParameters.to_cpp_default() + ), ) if not calc_compensations: compensations = None @@ -1268,7 +1276,7 @@ def forward( flatten_ids: Tensor, # [n_isects] absgrad: bool, ) -> Tuple[Tensor, Tensor]: - render_colors, render_alphas, last_ids = _make_lazy_cuda_func( + render_colors, render_alphas, last_ids = _make_lazy_device_func( "rasterize_to_pixels_3dgs_fwd" )( means2d, @@ -1334,7 +1342,7 @@ def backward( v_conics, v_colors, v_opacities, - ) = _make_lazy_cuda_func("rasterize_to_pixels_3dgs_bwd")( + ) = _make_lazy_device_func("rasterize_to_pixels_3dgs_bwd")( means2d, conics, colors, @@ -1412,7 +1420,7 @@ def forward( ) -> Tuple[Tensor, Tensor]: ut_params = ut_params.to_cpp() rs_type = rolling_shutter.to_cpp() - camera_model_type = _make_lazy_cuda_obj( + camera_model_type = _make_lazy_device_obj( f"CameraModelType.{camera_model.upper()}" ) ftheta_coeffs = ( @@ -1421,7 +1429,7 @@ def forward( else FThetaCameraDistortionParameters.to_cpp_default() ) - render_colors, render_alphas, last_ids = _make_lazy_cuda_func( + render_colors, render_alphas, last_ids = _make_lazy_device_func( "rasterize_to_pixels_from_world_3dgs_fwd" )( means, @@ -1510,7 +1518,7 @@ def backward( tile_size = ctx.tile_size ftheta_coeffs = ctx.ftheta_coeffs - (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_cuda_func( + (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_device_func( "rasterize_to_pixels_from_world_3dgs_bwd" )( means, @@ -1604,7 +1612,7 @@ def forward( camera_model != "ftheta" ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - camera_model_type = _make_lazy_cuda_obj( + camera_model_type = _make_lazy_device_obj( f"CameraModelType.{camera_model.upper()}" ) @@ -1618,7 +1626,7 @@ def forward( depths, conics, compensations, - ) = _make_lazy_cuda_func("projection_ewa_3dgs_packed_fwd")( + ) = _make_lazy_device_func("projection_ewa_3dgs_packed_fwd")( means, covars, # optional quats, # optional @@ -1702,7 +1710,7 @@ def backward( if v_compensations is not None: v_compensations = v_compensations.contiguous() - v_means, v_covars, v_quats, v_scales, v_viewmats = _make_lazy_cuda_func( + v_means, v_covars, v_quats, v_scales, v_viewmats = _make_lazy_device_func( "projection_ewa_3dgs_packed_bwd" )( means, @@ -1806,7 +1814,7 @@ class _SphericalHarmonics(torch.autograd.Function): def forward( ctx, sh_degree: int, dirs: Tensor, coeffs: Tensor, masks: Tensor ) -> Tensor: - colors = _make_lazy_cuda_func("spherical_harmonics_fwd")( + colors = _make_lazy_device_func("spherical_harmonics_fwd")( sh_degree, dirs, coeffs, masks ) ctx.save_for_backward(dirs, coeffs, masks) @@ -1820,7 +1828,7 @@ def backward(ctx, v_colors: Tensor): sh_degree = ctx.sh_degree num_bases = ctx.num_bases compute_v_dirs = ctx.needs_input_grad[1] - v_coeffs, v_dirs = _make_lazy_cuda_func("spherical_harmonics_bwd")( + v_coeffs, v_dirs = _make_lazy_device_func("spherical_harmonics_bwd")( num_bases, sh_degree, dirs, @@ -1959,7 +1967,7 @@ def forward( far_plane: float, radius_clip: float, ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - radii, means2d, depths, ray_transforms, normals = _make_lazy_cuda_func( + radii, means2d, depths, ray_transforms, normals = _make_lazy_device_func( "projection_2dgs_fused_fwd" )( means, @@ -2005,7 +2013,7 @@ def backward(ctx, v_radii, v_means2d, v_depths, v_ray_transforms, v_normals): width = ctx.width height = ctx.height eps2d = ctx.eps2d - v_means, v_quats, v_scales, v_viewmats = _make_lazy_cuda_func( + v_means, v_quats, v_scales, v_viewmats = _make_lazy_device_func( "projection_2dgs_fused_bwd" )( means, @@ -2076,7 +2084,7 @@ def forward( depths, ray_transforms, normals, - ) = _make_lazy_cuda_func("projection_2dgs_packed_fwd")( + ) = _make_lazy_device_func("projection_2dgs_packed_fwd")( means, quats, scales, @@ -2141,7 +2149,7 @@ def backward( height = ctx.height sparse_grad = ctx.sparse_grad - v_means, v_quats, v_scales, v_viewmats = _make_lazy_cuda_func( + v_means, v_quats, v_scales, v_viewmats = _make_lazy_device_func( "projection_2dgs_packed_bwd" )( means, @@ -2293,7 +2301,7 @@ def rasterize_to_pixels_2dgs( raise ValueError(f"Unsupported number of color channels: {channels}") if channels not in (1, 2, 3, 4, 8, 16, 32, 64, 128, 256, 512): padded_channels = (1 << (channels - 1).bit_length()) - channels - # Make sure the depth (last channel if present) remains in the last channel after padding (for depth distortion and median depth in CUDA kernel) + # Make sure the depth (last channel if present) remains in the last channel after padding (for depth distortion and median depth in device kernel) colors = torch.cat( [ colors[..., :-1], @@ -2421,7 +2429,7 @@ def rasterize_to_indices_in_range_2dgs( tile_width * tile_size >= image_width ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - out_gauss_ids, out_indices = _make_lazy_cuda_func("rasterize_to_indices_2dgs")( + out_gauss_ids, out_indices = _make_lazy_device_func("rasterize_to_indices_2dgs")( range_start, range_end, transmittances.contiguous(), @@ -2469,7 +2477,7 @@ def forward( render_median, last_ids, median_ids, - ) = _make_lazy_cuda_func("rasterize_to_pixels_2dgs_fwd")( + ) = _make_lazy_device_func("rasterize_to_pixels_2dgs_fwd")( means2d, ray_transforms, colors, @@ -2555,7 +2563,7 @@ def backward( v_opacities, v_normals, v_densify, - ) = _make_lazy_cuda_func("rasterize_to_pixels_2dgs_bwd")( + ) = _make_lazy_device_func("rasterize_to_pixels_2dgs_bwd")( means2d, ray_transforms, colors, @@ -2580,7 +2588,7 @@ def backward( v_render_median.contiguous(), absgrad, ) - torch.cuda.synchronize() + torch_acc.synchronize() if absgrad: means2d.absgrad = v_means2d_abs diff --git a/gsplat/cuda/csrc/third_party/glm b/gsplat/cuda/csrc/third_party/glm index 33b4a621..2d4c4b4d 160000 --- a/gsplat/cuda/csrc/third_party/glm +++ b/gsplat/cuda/csrc/third_party/glm @@ -1 +1 @@ -Subproject commit 33b4a621a697a305bc3a7610d290677b96beb181 +Subproject commit 2d4c4b4dd31fde06cfffad7915c2b3006402322f diff --git a/gsplat/distributed.py b/gsplat/distributed.py index cab559df..9b02bb52 100644 --- a/gsplat/distributed.py +++ b/gsplat/distributed.py @@ -6,6 +6,17 @@ import torch.distributed.nn.functional as distF from torch import Tensor +from . import torch_acc, BACKEND + + +def _get_distributed_backend(): + if BACKEND == "sycl": + return "xccl" + elif BACKEND == "cuda": + return "nccl" + else: + return "gloo" + def all_gather_int32( world_size: int, value: Union[int, Tensor], device: Optional[torch.device] = None @@ -30,13 +41,17 @@ def all_gather_int32( if world_size == 1: return [value] - # move to CUDA + # move to device if isinstance(value, int): assert device is not None, "device is required for scalar input" value_tensor = torch.tensor(value, dtype=torch.int, device=device) else: value_tensor = value - assert value_tensor.is_cuda, "value should be on CUDA" + + if BACKEND == "cuda": + assert value_tensor.is_cuda, "value should be on CUDA" + elif BACKEND == "sycl": + assert value_tensor.is_xpu, "value should be on XPU" # gather collected = torch.empty( @@ -82,7 +97,7 @@ def all_to_all_int32( if any(isinstance(v, int) for v in values): assert device is not None, "device is required for scalar input" - # move to CUDA + # move to device values_tensor = [ (torch.tensor(v, dtype=torch.int, device=device) if isinstance(v, int) else v) for v in values @@ -283,9 +298,10 @@ def _distributed_worker( print("Distributed worker: %d / %d" % (world_rank + 1, world_size)) distributed = world_size > 1 if distributed: - torch.cuda.set_device(local_rank) + torch_acc.set_device(local_rank) + torch.distributed.init_process_group( - backend="nccl", world_size=world_size, rank=world_rank + backend=_get_distributed_backend(), world_size=world_size, rank=world_rank ) # Dump collection that participates all ranks. # This initializes the communicator required by `batch_isend_irecv`. @@ -319,7 +335,6 @@ def fn(local_rank: int, world_rank: int, world_size: int, args: Any) -> None: cli(fn, None, verbose=True) ``` """ - assert torch.cuda.is_available(), "CUDA device is required!" if "OMPI_COMM_WORLD_SIZE" in os.environ: # multi-node local_rank = int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"]) world_size = int(os.environ["OMPI_COMM_WORLD_SIZE"]) # dist.get_world_size() @@ -328,7 +343,7 @@ def fn(local_rank: int, world_rank: int, world_size: int, args: Any) -> None: world_rank, world_size, fn, args, local_rank, verbose ) - world_size = torch.cuda.device_count() + world_size = torch_acc.device_count() distributed = world_size > 1 if distributed: diff --git a/gsplat/optimizers/selective_adam.py b/gsplat/optimizers/selective_adam.py index 02f66281..b4f284b6 100644 --- a/gsplat/optimizers/selective_adam.py +++ b/gsplat/optimizers/selective_adam.py @@ -1,6 +1,6 @@ import torch -from ..cuda._wrapper import adam +from .._wrapper import adam class SelectiveAdam(torch.optim.Adam): diff --git a/gsplat/profile.py b/gsplat/profile.py index 669d363f..8d31af65 100644 --- a/gsplat/profile.py +++ b/gsplat/profile.py @@ -4,6 +4,7 @@ from typing import Callable, Optional import torch +from gsplat import torch_acc profiler = {} @@ -36,12 +37,12 @@ def __init__(self, name: str = "unnamed"): def __enter__(self): if self.enabled: - torch.cuda.synchronize() + torch_acc.synchronize() self.start_time = time.perf_counter() def __exit__(self, exc_type, exc_val, exc_tb): if self.enabled: - torch.cuda.synchronize() + torch_acc.synchronize() end_time = time.perf_counter() total_time = end_time - self.start_time if self.name not in profiler: diff --git a/gsplat/relocation.py b/gsplat/relocation.py index 8abd9aee..0a6600d5 100644 --- a/gsplat/relocation.py +++ b/gsplat/relocation.py @@ -4,7 +4,8 @@ import torch from torch import Tensor -from .cuda._wrapper import _make_lazy_cuda_func +from . import BACKEND +from ._wrapper import _make_lazy_device_func def compute_relocation( @@ -43,7 +44,7 @@ def compute_relocation( ratios.clamp_(min=1, max=n_max) ratios = ratios.int().contiguous() - new_opacities, new_scales = _make_lazy_cuda_func("relocation")( + new_opacities, new_scales = _make_lazy_device_func("relocation")( opacities, scales, ratios, binoms, n_max ) return new_opacities, new_scales diff --git a/gsplat/rendering.py b/gsplat/rendering.py index d103e952..875cecc5 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -7,10 +7,9 @@ from torch import Tensor from typing_extensions import Literal -from .cuda._wrapper import ( +from . import BACKEND +from ._wrapper import ( RollingShutterType, - FThetaCameraDistortionParameters, - FThetaPolynomialType, fully_fused_projection, fully_fused_projection_2dgs, fully_fused_projection_with_ut, @@ -73,7 +72,7 @@ def _compute_view_dirs_packed( avg_means_per_camera = nnz / (B * C) split_batch_camera_ops = ( avg_means_per_camera > 10000 - and campos_flat.is_cuda + and not campos_flat.is_cpu and campos_flat.requires_grad ) @@ -138,7 +137,7 @@ def rasterization( radial_coeffs: Optional[Tensor] = None, # [..., C, 6] or [..., C, 4] tangential_coeffs: Optional[Tensor] = None, # [..., C, 2] thin_prism_coeffs: Optional[Tensor] = None, # [..., C, 4] - ftheta_coeffs: Optional[FThetaCameraDistortionParameters] = None, + ftheta_coeffs=None, # rolling shutter rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] @@ -623,7 +622,7 @@ def reshape_view(C: int, world_view: torch.Tensor, N_world: list) -> torch.Tenso (radii,) = all_to_all_tensor_list( world_size, [radii], cnts, output_splits=collected_splits ) - (means2d, depths, conics, opacities, colors) = all_to_all_tensor_list( + means2d, depths, conics, opacities, colors = all_to_all_tensor_list( world_size, [means2d, depths, conics, opacities, colors], cnts, @@ -651,7 +650,7 @@ def reshape_view(C: int, world_view: torch.Tensor, N_world: list) -> torch.Tenso gaussian_ids = gaussian_ids + offsets # all to all communication across all ranks. - (camera_ids, gaussian_ids) = all_to_all_tensor_list( + camera_ids, gaussian_ids = all_to_all_tensor_list( world_size, [camera_ids, gaussian_ids], cnts, @@ -675,7 +674,7 @@ def reshape_view(C: int, world_view: torch.Tensor, N_world: list) -> torch.Tenso ) radii = reshape_view(C, radii, N_world) - (means2d, depths, conics, opacities, colors) = all_to_all_tensor_list( + means2d, depths, conics, opacities, colors = all_to_all_tensor_list( world_size, [ means2d.flatten(0, 1), @@ -878,7 +877,7 @@ def _rasterization( .. note:: This function still relies on gsplat's CUDA backend for some computation, but the - entire differentiable graph is on of PyTorch (and nerfacc) so could use Pytorch's + entire differentiable graph is on PyTorch (and nerfacc) so could use Pytorch's autograd for backpropagation. .. note:: @@ -889,7 +888,7 @@ def _rasterization( Compared to rasterization(), this function does not support some arguments such as `packed`, `sparse_grad` and `absgrad`. """ - from gsplat.cuda._torch_impl import ( + from gsplat._torch_impl import ( _fully_fused_projection, _quat_scale_to_covar_preci, _rasterize_to_pixels, @@ -1561,7 +1560,7 @@ def rasterization_2dgs( image_ids = None densify = torch.zeros_like( - means2d, dtype=means.dtype, requires_grad=True, device="cuda" + means2d, dtype=means.dtype, requires_grad=True, device=means2d.device ) # Identify intersecting tiles tile_width = math.ceil(width / float(tile_size)) diff --git a/gsplat/strategy/default.py b/gsplat/strategy/default.py index b19b73be..3b112c31 100644 --- a/gsplat/strategy/default.py +++ b/gsplat/strategy/default.py @@ -6,6 +6,7 @@ from .base import Strategy from .ops import duplicate, remove, reset_opa, split +from .. import torch_acc @dataclass @@ -190,7 +191,7 @@ def step_post_backward( state["count"].zero_() if self.refine_scale2d_stop_iter > 0: state["radii"].zero_() - torch.cuda.empty_cache() + torch_acc.empty_cache() if step % self.reset_every == 0 and step > 0: reset_opa( diff --git a/gsplat/strategy/mcmc.py b/gsplat/strategy/mcmc.py index c07e1737..98689bf9 100644 --- a/gsplat/strategy/mcmc.py +++ b/gsplat/strategy/mcmc.py @@ -5,6 +5,7 @@ import torch from torch import Tensor +from gsplat import torch_acc from .base import Strategy from .ops import inject_noise_to_position, relocate, sample_add @@ -137,7 +138,7 @@ def step_post_backward( f"Now having {len(params['means'])} GSs." ) - torch.cuda.empty_cache() + torch_acc.empty_cache() # add noise to GSs inject_noise_to_position( diff --git a/gsplat/strategy/ops.py b/gsplat/strategy/ops.py index 83c90a25..9f6892d8 100644 --- a/gsplat/strategy/ops.py +++ b/gsplat/strategy/ops.py @@ -5,9 +5,9 @@ import torch.nn.functional as F from torch import Tensor -from gsplat import quat_scale_to_covar_preci -from gsplat.relocation import compute_relocation -from gsplat.utils import normalized_quat_to_rotmat +from .._wrapper import quat_scale_to_covar_preci +from ..relocation import compute_relocation +from ..utils import normalized_quat_to_rotmat @torch.no_grad() diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt new file mode 100644 index 00000000..4644f18e --- /dev/null +++ b/gsplat/sycl/CMakeLists.txt @@ -0,0 +1,133 @@ +cmake_minimum_required(VERSION 3.23...4.0) # Need min 3.23 on Windows + +set(CMAKE_C_COMPILER icx) +set(CMAKE_CXX_COMPILER icx) + +project(gsplat_sycl) + +if (NOT CMAKE_BUILD_TYPE AND NOT CMAKE_CONFIGURATION_TYPES) + set(CMAKE_BUILD_TYPE Release CACHE STRING "Choose the type of build." FORCE) + set_property(CACHE CMAKE_BUILD_TYPE PROPERTY STRINGS "Debug" "Release" "MinSizeRel" "RelWithDebInfo") +endif() + +if (NOT SYCL_AOT_TARGETS) + set (SYCL_AOT_TARGETS "spir64" CACHE STRING "Comma separated list of SYCL targets for ahead of time compilation. See https://github.com/intel/llvm/blob/sycl/sycl/doc/UsersManual.md for a full list.") +endif() + +find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED) + +execute_process( + COMMAND "${Python_EXECUTABLE}" -c " +import sys, os, glob +try: + import torch + from torch.utils import cpp_extension + import pybind11 +except ImportError: + sys.exit(1) + +lib_paths = cpp_extension.library_paths(True) +lib_name = 'libtorch_python.so*' if sys.platform != 'win32' else 'torch_python.lib' +base_path = os.path.join(lib_paths[0], lib_name) +matches = glob.glob(base_path) +torch_python_lib = matches[0] if matches else '' + +print(f'{torch.utils.cmake_prefix_path};{torch_python_lib};{pybind11.get_cmake_dir()}') +" + OUTPUT_VARIABLE PYTHON_CONFIG_LIST + RESULT_VARIABLE PYTHON_CONFIG_RESULT + OUTPUT_STRIP_TRAILING_WHITESPACE +) + +if(NOT PYTHON_CONFIG_RESULT EQUAL 0) + message(FATAL_ERROR "Failed to retrieve configuration from Python. Ensure torch and pybind11 are installed.") +endif() + +list(LENGTH PYTHON_CONFIG_LIST LIST_LEN) +if(LIST_LEN LESS 3) + message(FATAL_ERROR "Python script returned incomplete information.") +endif() +list(GET PYTHON_CONFIG_LIST 0 Torch_DIR_From_Python) +list(GET PYTHON_CONFIG_LIST 1 TORCH_PYTHON_LIB) +list(GET PYTHON_CONFIG_LIST 2 PYBIND11_CMAKE_DIR) + +if(Torch_DIR_From_Python AND IS_DIRECTORY "${Torch_DIR_From_Python}") + set(Torch_DIR ${Torch_DIR_From_Python}) + message(STATUS "Found Torch CMake directory via Python: ${Torch_DIR}") + find_package(Torch REQUIRED HINTS ${Torch_DIR_From_Python}) +else() + message(FATAL_ERROR "Could not find Torch via Python introspection.") +endif() + +if (NOT EXISTS "${TORCH_PYTHON_LIB}") + message(FATAL_ERROR "Could not find torch_python library.") +else() + message(STATUS "Found torch_python library at: ${TORCH_PYTHON_LIB}") +endif() + +message(STATUS "Found pybind11 in ${PYBIND11_CMAKE_DIR}") +find_package(pybind11 CONFIG REQUIRED HINTS ${PYBIND11_CMAKE_DIR} "${Python_SITELIB}/pybind11/share/cmake/pybind11") + +set(SYCL_SOURCES + ext.cpp + src/adam.cpp + src/intersect_offset.cpp + src/intersect_tile.cpp + src/null.cpp + src/projection_2dgs_fused_bwd.cpp + src/projection_2dgs_fused_fwd.cpp + src/projection_2dgs_packed_bwd.cpp + src/projection_2dgs_packed_fwd.cpp + src/projection_ewa_3dgs_fused_bwd.cpp + src/projection_ewa_3dgs_fused_fwd.cpp + src/projection_ewa_3dgs_packed_bwd.cpp + src/projection_ewa_3dgs_packed_fwd.cpp + src/projection_ewa_simple_bwd.cpp + src/projection_ewa_simple_fwd.cpp + src/projection_ut_3dgs_fused.cpp + src/quat_scale_to_covar_preci_bwd.cpp + src/quat_scale_to_covar_preci_fwd.cpp + src/rasterize_to_indices_2dgs.cpp + src/rasterize_to_indices_3dgs.cpp + src/rasterize_to_pixels_2dgs_bwd.cpp + src/rasterize_to_pixels_2dgs_fwd.cpp + src/rasterize_to_pixels_3dgs_bwd.cpp + src/rasterize_to_pixels_3dgs_fwd.cpp + src/rasterize_to_pixels_from_world_3dgs_bwd.cpp + src/rasterize_to_pixels_from_world_3dgs_fwd.cpp + src/relocation.cpp + src/spherical_harmonics_bwd.cpp + src/spherical_harmonics_fwd.cpp +) + +set(SYCL_MODULE_NAME gsplat_sycl_kernels) + +pybind11_add_module(${SYCL_MODULE_NAME} MODULE ${SYCL_SOURCES}) + +target_include_directories( ${SYCL_MODULE_NAME} SYSTEM PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR}/../cuda/csrc/third_party/glm +) + +target_compile_definitions(${SYCL_MODULE_NAME} PRIVATE + TORCH_EXTENSION_NAME=${SYCL_MODULE_NAME} +) + +target_compile_options(${SYCL_MODULE_NAME} PRIVATE -fsycl) +target_compile_features(${SYCL_MODULE_NAME} PUBLIC cxx_std_17) + +target_link_options(${SYCL_MODULE_NAME} PRIVATE -fsycl -fsycl-targets=${SYCL_AOT_TARGETS}) +target_link_libraries(${SYCL_MODULE_NAME} PRIVATE torch ${TORCH_PYTHON_LIB}) + + +# Fix for icx: error: '-MP' is not supported with offloading enabled +if (WIN32 AND NOT UNIX) + get_target_property(CURRENT_OPTIONS ${SYCL_MODULE_NAME} COMPILE_OPTIONS) + string(REPLACE "/MP" "" MODIFIED_OPTIONS "${CURRENT_OPTIONS}") + set_target_properties(${SYCL_MODULE_NAME} PROPERTIES COMPILE_OPTIONS "${MODIFIED_OPTIONS}") +endif () +if (UNIX AND NOT APPLE) + # Find libtorch_xpu and libsycl at runtime in Python environment + set_target_properties(${SYCL_MODULE_NAME} PROPERTIES INSTALL_RPATH + "$ORIGIN/../../torch/lib/;$ORIGIN/../../../../") +endif() \ No newline at end of file diff --git a/gsplat/sycl/__init__.py b/gsplat/sycl/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/gsplat/sycl/_backend.py b/gsplat/sycl/_backend.py new file mode 100644 index 00000000..abe4e60e --- /dev/null +++ b/gsplat/sycl/_backend.py @@ -0,0 +1,26 @@ +_C = None + +import os + +import gsplat + +if os.name == "nt": + import torch + from sysconfig import get_path + + dllpath = [ + os.add_dll_directory(torch.__path__[0] + "/lib"), # for torch libs + os.add_dll_directory(get_path("data") + "/Library/bin"), # for sycl libs + ] + +try: + # Try to import the compiled module (via setup.py or pre-built .so) + from gsplat import gsplat_sycl_kernels as _C +except ImportError: + raise ImportError("Unable to find compiled sycl kernels package") + +if os.name == "nt": + for dp in dllpath: + dp.close() + +__all__ = ["_C"] diff --git a/gsplat/sycl/ext.cpp b/gsplat/sycl/ext.cpp new file mode 100644 index 00000000..74346ca0 --- /dev/null +++ b/gsplat/sycl/ext.cpp @@ -0,0 +1,150 @@ +#include + +#include "Cameras.h" +#include "Ops.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + + py::enum_(m, "CameraModelType") + .value("PINHOLE", gsplat::xpu::CameraModelType::PINHOLE) + .value("ORTHO", gsplat::xpu::CameraModelType::ORTHO) + .value("FISHEYE", gsplat::xpu::CameraModelType::FISHEYE) + .value("FTHETA", gsplat::xpu::CameraModelType::FTHETA) + .export_values(); + + m.def("null", &gsplat::xpu::null); + + m.def( + "quat_scale_to_covar_preci_fwd", + &gsplat::xpu::quat_scale_to_covar_preci_fwd + ); + m.def( + "quat_scale_to_covar_preci_bwd", + &gsplat::xpu::quat_scale_to_covar_preci_bwd + ); + + m.def("spherical_harmonics_fwd", &gsplat::xpu::spherical_harmonics_fwd); + m.def("spherical_harmonics_bwd", &gsplat::xpu::spherical_harmonics_bwd); + + m.def("adam", &gsplat::xpu::adam); + m.def("relocation", &gsplat::xpu::relocation); + + m.def("intersect_tile", &gsplat::xpu::intersect_tile); + m.def("intersect_offset", &gsplat::xpu::intersect_offset); + + m.def("projection_ewa_simple_fwd", &gsplat::xpu::projection_ewa_simple_fwd); + m.def("projection_ewa_simple_bwd", &gsplat::xpu::projection_ewa_simple_bwd); + m.def( + "projection_ewa_3dgs_fused_fwd", + &gsplat::xpu::projection_ewa_3dgs_fused_fwd + ); + m.def( + "projection_ewa_3dgs_fused_bwd", + &gsplat::xpu::projection_ewa_3dgs_fused_bwd + ); + m.def( + "projection_ewa_3dgs_packed_fwd", + &gsplat::xpu::projection_ewa_3dgs_packed_fwd + ); + m.def( + "projection_ewa_3dgs_packed_bwd", + &gsplat::xpu::projection_ewa_3dgs_packed_bwd + ); + + m.def( + "rasterize_to_pixels_3dgs_fwd", + &gsplat::xpu::rasterize_to_pixels_3dgs_fwd + ); + m.def( + "rasterize_to_pixels_3dgs_bwd", + &gsplat::xpu::rasterize_to_pixels_3dgs_bwd + ); + m.def("rasterize_to_indices_3dgs", &gsplat::xpu::rasterize_to_indices_3dgs); + + m.def("projection_2dgs_fused_fwd", &gsplat::xpu::projection_2dgs_fused_fwd); + m.def("projection_2dgs_fused_bwd", &gsplat::xpu::projection_2dgs_fused_bwd); + m.def( + "projection_2dgs_packed_fwd", &gsplat::xpu::projection_2dgs_packed_fwd + ); + m.def( + "projection_2dgs_packed_bwd", &gsplat::xpu::projection_2dgs_packed_bwd + ); + + m.def( + "rasterize_to_pixels_2dgs_fwd", + &gsplat::xpu::rasterize_to_pixels_2dgs_fwd + ); + m.def( + "rasterize_to_pixels_2dgs_bwd", + &gsplat::xpu::rasterize_to_pixels_2dgs_bwd + ); + m.def("rasterize_to_indices_2dgs", &gsplat::xpu::rasterize_to_indices_2dgs); + + m.def("projection_ut_3dgs_fused", &gsplat::xpu::projection_ut_3dgs_fused); + m.def( + "rasterize_to_pixels_from_world_3dgs_fwd", + &gsplat::xpu::rasterize_to_pixels_from_world_3dgs_fwd + ); + m.def( + "rasterize_to_pixels_from_world_3dgs_bwd", + &gsplat::xpu::rasterize_to_pixels_from_world_3dgs_bwd + ); + + // Cameras from 3DGUT + py::enum_(m, "ShutterType") + .value("ROLLING_TOP_TO_BOTTOM", ShutterType::ROLLING_TOP_TO_BOTTOM) + .value("ROLLING_LEFT_TO_RIGHT", ShutterType::ROLLING_LEFT_TO_RIGHT) + .value("ROLLING_BOTTOM_TO_TOP", ShutterType::ROLLING_BOTTOM_TO_TOP) + .value("ROLLING_RIGHT_TO_LEFT", ShutterType::ROLLING_RIGHT_TO_LEFT) + .value("GLOBAL", ShutterType::GLOBAL) + .export_values(); + + py::class_(m, "UnscentedTransformParameters") + .def(py::init<>()) + .def_readwrite("alpha", &UnscentedTransformParameters::alpha) + .def_readwrite("beta", &UnscentedTransformParameters::beta) + .def_readwrite("kappa", &UnscentedTransformParameters::kappa) + .def_readwrite( + "in_image_margin_factor", + &UnscentedTransformParameters::in_image_margin_factor + ) + .def_readwrite( + "require_all_sigma_points_valid", + &UnscentedTransformParameters::require_all_sigma_points_valid + ); + + // FTheta Camera support + py::enum_( + m, "FThetaPolynomialType" + ) + .value( + "PIXELDIST_TO_ANGLE", + FThetaCameraDistortionParameters::PolynomialType::PIXELDIST_TO_ANGLE + ) + .value( + "ANGLE_TO_PIXELDIST", + FThetaCameraDistortionParameters::PolynomialType::ANGLE_TO_PIXELDIST + ) + .export_values(); + py::class_( + m, "FThetaCameraDistortionParameters" + ) + .def(py::init<>()) + .def_readwrite( + "reference_poly", &FThetaCameraDistortionParameters::reference_poly + ) + .def_readwrite( + "pixeldist_to_angle_poly", + &FThetaCameraDistortionParameters::pixeldist_to_angle_poly + ) + .def_readwrite( + "angle_to_pixeldist_poly", + &FThetaCameraDistortionParameters::angle_to_pixeldist_poly + ) + .def_readwrite( + "max_angle", &FThetaCameraDistortionParameters::max_angle + ) + .def_readwrite( + "linear_cde", &FThetaCameraDistortionParameters::linear_cde + ); +} \ No newline at end of file diff --git a/gsplat/sycl/include/Cameras.h b/gsplat/sycl/include/Cameras.h new file mode 100644 index 00000000..bd3d38ed --- /dev/null +++ b/gsplat/sycl/include/Cameras.h @@ -0,0 +1,60 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +// --------------------------------------------------------------------------------------------- + +// Camera-specific types (camera model parameters and returns) + +enum class ShutterType { + ROLLING_TOP_TO_BOTTOM, + ROLLING_LEFT_TO_RIGHT, + ROLLING_BOTTOM_TO_TOP, + ROLLING_RIGHT_TO_LEFT, + GLOBAL +}; + +// --------------------------------------------------------------------------------------------- + +// Gaussian-specific types +struct UnscentedTransformParameters { + // See Gustafsson and Hendeby 2012 for sigma point parameterization - this + // default parameter choice is based on + // + // - "The unscented Kalman filter for nonlinear estimation" - Wan and van + // der Merwe 2000 + float alpha = 0.1; + float beta = 2.f; + float kappa = 0.f; + + // Parameters controlling validity of the unscented transform results + float in_image_margin_factor = + 0.1f; // 10% out of bounds margin is acceptable for "valid" projection + // state + bool require_all_sigma_points_valid = + false; // true: all sigma points must be valid to mark a projection as + // "valid" false: a single valid sigma point is sufficient to + // mark a projection as "valid" +}; + +// FTheta Camera Support +struct FThetaCameraDistortionParameters { + static constexpr size_t PolynomialDegree = 6; + enum class PolynomialType { + PIXELDIST_TO_ANGLE, + ANGLE_TO_PIXELDIST, + }; + PolynomialType reference_poly; + std::array + pixeldist_to_angle_poly; // backward polynomial + std::array + angle_to_pixeldist_poly; // forward polynomial + float max_angle; + std::array linear_cde; +}; \ No newline at end of file diff --git a/gsplat/sycl/include/Common.h b/gsplat/sycl/include/Common.h new file mode 100644 index 00000000..87b222aa --- /dev/null +++ b/gsplat/sycl/include/Common.h @@ -0,0 +1,43 @@ +#pragma once + +#include +#include +#include + +namespace gsplat::xpu { + +// +// Some Macros. +// +#define CHECK_XPU(x) TORCH_CHECK(x.is_xpu(), #x " must be a XPU tensor") +#define CHECK_DEVICE(x, y) \ + TORCH_CHECK( \ + x.device() == y.device(), #x " must be on device " + y.device().str() \ + ) +#define CHECK_CONTIGUOUS(x) \ + TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) \ + { \ + CHECK_XPU(x); \ + CHECK_CONTIGUOUS(x); \ + } +#define CHECK_INPUT2(x, y) \ + { \ + CHECK_DEVICE(x, y); \ + CHECK_CONTIGUOUS(x); \ + } +#define DEVICE_GUARD(_ten) const c10::DeviceGuard device_guard(_ten.device()); + +// +// Legacy Camera Types +// +enum CameraModelType { + PINHOLE = 0, + ORTHO = 1, + FISHEYE = 2, + FTHETA = 3, +}; + +#define GSPLAT_N_THREADS 256 +#define N_THREADS_PACKED 256 +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/Ops.h b/gsplat/sycl/include/Ops.h new file mode 100644 index 00000000..80ec7e8a --- /dev/null +++ b/gsplat/sycl/include/Ops.h @@ -0,0 +1,574 @@ +// A collection of operators for gsplat +#pragma once + +#include "Cameras.h" +#include "Common.h" +#include "types.hpp" +#include +#include + +namespace gsplat::xpu { + +// null operator for tutorial. Does nothing. +at::Tensor null(const at::Tensor input); + +// Project 3D gaussians (in camera space) to 2D image planes with EWA splatting. +std::tuple projection_ewa_simple_fwd( + const at::Tensor means, // [..., C, N, 3] + const at::Tensor covars, // [..., C, N, 3, 3] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t width, + const uint32_t height, + const CameraModelType camera_model +); +std::tuple projection_ewa_simple_bwd( + const at::Tensor means, // [..., C, N, 3] + const at::Tensor covars, // [..., C, N, 3, 3] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t width, + const uint32_t height, + const CameraModelType camera_model, + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_covars2d // [..., C, N, 2, 2] +); + +// Fuse the following operations: +// 1. compute covar from {quats, scales} +// 2. transform 3D gaussians from world space to camera space +// - w/ near far plane check +// 3. projection camera space 3D gaussians to 2D image planes with EWA +// splatting. +// - w/ minimum radius check +// 4. add a bit blurring to the 2D gaussians for anti-aliasing. +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_ewa_3dgs_fused_fwd( + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] optional + const at::optional quats, // [..., N, 4] optional + const at::optional scales, // [..., N, 3] optional + const at::optional opacities, // [..., N] optional + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip, + const bool calc_compensations, + const CameraModelType camera_model +); +std::tuple +projection_ewa_3dgs_fused_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] optional + const at::optional quats, // [..., N, 4] optional + const at::optional scales, // [..., N, 3] optional + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const CameraModelType camera_model, + // fwd outputs + const at::Tensor radii, // [..., C, N, 2] + const at::Tensor conics, // [..., C, N, 3] + const at::optional compensations, // [..., C, N] optional + // grad outputs + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_depths, // [..., C, N] + const at::Tensor v_conics, // [..., C, N, 3] + const at::optional v_compensations, // [..., C, N] optional + const bool viewmats_requires_grad +); + +// On top of fusing the operations like `projection_ewa_3dgs_fused_{fwd, bwd}`, +// The packed version compresses the [C, N, D] tensors (both intermidiate and +// output) into a jagged format [nnz, D], leveraging the sparsity of these +// tensors. +// +// This could lead to less memory usage than `_fused_{fwd, bwd}` if the level of +// sparsity is high, i.e., most of the gaussians are not in the camera frustum. +// But at the cost of slightly slower speed. +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_ewa_3dgs_packed_fwd( + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] optional + const at::optional quats, // [..., N, 4] optional + const at::optional scales, // [..., N, 3] optional + const at::optional opacities, // [..., N] optional + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip, + const bool calc_compensations, + const CameraModelType camera_model +); +std::tuple +projection_ewa_3dgs_packed_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] + const at::optional quats, // [..., N, 4] + const at::optional scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const CameraModelType camera_model, + // fwd outputs + const at::Tensor batch_ids, // [nnz] + const at::Tensor camera_ids, // [nnz] + const at::Tensor gaussian_ids, // [nnz] + const at::Tensor conics, // [nnz, 3] + const at::optional compensations, // [nnz] optional + // grad outputs + const at::Tensor v_means2d, // [nnz, 2] + const at::Tensor v_depths, // [nnz] + const at::Tensor v_conics, // [nnz, 3] + const at::optional v_compensations, // [nnz] optional + const bool viewmats_requires_grad, + const bool sparse_grad +); + +// Sphereical harmonics +at::Tensor spherical_harmonics_fwd( + const uint32_t degrees_to_use, + const at::Tensor dirs, // [..., 3] + const at::Tensor coeffs, // [..., K, 3] + const at::optional masks // [...] +); +std::tuple spherical_harmonics_bwd( + const uint32_t K, + const uint32_t degrees_to_use, + const at::Tensor dirs, // [..., 3] + const at::Tensor coeffs, // [..., K, 3] + const at::optional masks, // [...] + const at::Tensor v_colors, // [..., 3] + bool compute_v_dirs +); + +// Fused Adam that supports a valid mask to skip updating certain parameters. +// Note skipping is not equivalent with zeroing out the gradients, which will +// still update parameters with momentum. +void adam( + at::Tensor ¶m, // [..., D] + const at::Tensor ¶m_grad, // [..., D] + at::Tensor &exp_avg, // [..., D] + at::Tensor &exp_avg_sq, // [..., D] + const at::optional valid, // [...] + const float lr, + const float b1, + const float b2, + const float eps +); + +// GS Tile Intersection +std::tuple intersect_tile( + const at::Tensor means2d, // [..., C, N, 2] or [nnz, 2] + const at::Tensor radii, // [..., C, N, 2] or [nnz, 2] + const at::Tensor depths, // [..., C, N] or [nnz] + const at::optional image_ids, // [nnz] + const at::optional gaussian_ids, // [nnz] + const uint32_t I, + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + const bool sort, + const bool segmented +); +at::Tensor intersect_offset( + const at::Tensor isect_ids, // [n_isects] + const uint32_t I, + const uint32_t tile_width, + const uint32_t tile_height +); + +// Compute Covariance and Precision Matrices from Quaternion and Scale +std::tuple quat_scale_to_covar_preci_fwd( + const at::Tensor quats, // [..., 4] + const at::Tensor scales, // [..., 3] + const bool compute_covar, + const bool compute_preci, + const bool triu +); +std::tuple quat_scale_to_covar_preci_bwd( + const at::Tensor quats, // [..., 4] + const at::Tensor scales, // [..., 3] + const bool triu, + const at::optional v_covars, // [..., 3, 3] or [..., 6] + const at::optional v_precis // [..., 3, 3] or [..., 6] +); + +// Rasterize 3D Gaussian to pixels +std::tuple rasterize_to_pixels_3dgs_fwd( + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] or [nnz, 2] + const at::Tensor conics, // [..., N, 3] or [nnz, 3] + const at::Tensor colors, // [..., N, channels] or [nnz, channels] + const at::Tensor opacities, // [..., N] or [nnz] + const at::optional backgrounds, // [..., channels] + const at::optional masks, // [..., tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +); +std::tuple +rasterize_to_pixels_3dgs_bwd( + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] or [nnz, 2] + const at::Tensor conics, // [..., N, 3] or [nnz, 3] + const at::Tensor colors, // [..., N, 3] or [nnz, 3] + const at::Tensor opacities, // [..., N] or [nnz] + const at::optional backgrounds, // [..., 3] + const at::optional masks, // [..., tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids, // [n_isects] + // forward outputs + const at::Tensor render_alphas, // [..., image_height, image_width, 1] + const at::Tensor last_ids, // [..., image_height, image_width] + // gradients of outputs + const at::Tensor v_render_colors, // [..., image_height, image_width, 3] + const at::Tensor v_render_alphas, // [..., image_height, image_width, 1] + // options + bool absgrad +); + +// Rasterize 3D Gaussian, but only return the indices of gaussians and pixels. +std::tuple rasterize_to_indices_3dgs( + const uint32_t range_start, + const uint32_t range_end, // iteration steps + const at::Tensor transmittances, // [..., image_height, image_width] + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] + const at::Tensor conics, // [..., N, 3] + const at::Tensor opacities, // [..., N] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +); + +// Relocate some Gaussians in the Densification Process. +// Equation (9) in "3D Gaussian Splatting as Markov Chain Monte Carlo" +std::tuple relocation( + at::Tensor opacities, // [N] + at::Tensor scales, // [N, 3] + at::Tensor ratios, // [N] + at::Tensor binoms, // [n_max, n_max] + const int n_max +); + +// Projection for 2DGS +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_2dgs_fused_fwd( + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip +); +std::tuple +projection_2dgs_fused_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + // fwd outputs + const at::Tensor radii, // [..., C, N, 2] + const at::Tensor ray_transforms, // [..., C, N, 3, 3] + // grad outputs + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_depths, // [..., C, N] + const at::Tensor v_normals, // [..., C, N, 3] + const at::Tensor v_ray_transforms, // [..., C, N, 3, 3] + const bool viewmats_requires_grad +); + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_2dgs_packed_fwd( + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float near_plane, + const float far_plane, + const float radius_clip +); +std::tuple +projection_2dgs_packed_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + // fwd outputs + const at::Tensor batch_ids, // [nnz] + const at::Tensor camera_ids, // [nnz] + const at::Tensor gaussian_ids, // [nnz] + const at::Tensor ray_transforms, // [nnz, 3, 3] + // grad outputs + const at::Tensor v_means2d, // [nnz, 2] + const at::Tensor v_depths, // [nnz] + const at::Tensor v_ray_transforms, // [nnz, 3, 3] + const at::Tensor v_normals, // [nnz, 3] + const bool viewmats_requires_grad, + const bool sparse_grad +); + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +rasterize_to_pixels_2dgs_fwd( + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] or [nnz, 2] + const at::Tensor ray_transforms, // [..., N, 3, 3] or [nnz, 3, 3] + const at::Tensor colors, // [..., N, channels] or [nnz, channels] + const at::Tensor opacities, // [..., N] or [nnz] + const at::Tensor normals, // [..., N, 3] or [nnz, 3] + const at::optional backgrounds, // [..., channels] + const at::optional masks, // [..., tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +); +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +rasterize_to_pixels_2dgs_bwd( + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] or [nnz, 2] + const at::Tensor ray_transforms, // [..., N, 3, 3] or [nnz, 3, 3] + const at::Tensor colors, // [..., N, 3] or [nnz, 3] + const at::Tensor opacities, // [..., N] or [nnz] + const at::Tensor normals, // [..., N, 3] or [nnz, 3] + const at::Tensor densify, + const at::optional backgrounds, // [..., 3] + const at::optional masks, // [..., tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // ray_crossions + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids, // [n_isects] + // forward outputs + const at::Tensor + render_colors, // [..., image_height, image_width, COLOR_DIM] + const at::Tensor render_alphas, // [..., image_height, image_width, 1] + const at::Tensor last_ids, // [..., image_height, image_width] + const at::Tensor median_ids, // [..., image_height, image_width] + // gradients of outputs + const at::Tensor v_render_colors, // [..., image_height, image_width, 3] + const at::Tensor v_render_alphas, // [..., image_height, image_width, 1] + const at::Tensor v_render_normals, // [..., image_height, image_width, 3] + const at::Tensor v_render_distort, // [..., image_height, image_width, 1] + const at::Tensor v_render_median, // [..., image_height, image_width, 1] + // options + bool absgrad +); + +std::tuple rasterize_to_indices_2dgs( + const uint32_t range_start, + const uint32_t range_end, // iteration steps + const at::Tensor transmittances, // [..., image_height, image_width] + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] + const at::Tensor ray_transforms, // [..., N, 3, 3] + const at::Tensor opacities, // [..., N] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +); + +// Use uncented transform to project 3D gaussians to 2D. (none differentiable) +// https://arxiv.org/abs/2412.12507 +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_ut_3dgs_fused( + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::optional opacities, // [..., N] optional + const at::Tensor viewmats0, // [..., C, 4, 4] + const at::optional + viewmats1, // [..., C, 4, 4] optional for rolling shutter + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip, + const bool calc_compensations, + const CameraModelType camera_model, + // uncented transform + const UnscentedTransformParameters ut_params, + ShutterType rs_type, + const at::optional + radial_coeffs, // [..., C, 6] or [..., C, 4] optional + const at::optional tangential_coeffs, // [..., C, 2] optional + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs // shared parameters for all cameras +); + +std::tuple +rasterize_to_pixels_from_world_3dgs_fwd( + // Gaussian parameters + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor colors, // [..., C, N, channels] or [nnz, channels] + const at::Tensor opacities, // [..., C, N] or [nnz] + const at::optional backgrounds, // [..., C, channels] + const at::optional masks, // [..., C, tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // camera + const at::Tensor viewmats0, // [..., C, 4, 4] + const at::optional + viewmats1, // [..., C, 4, 4] optional for rolling shutter + const at::Tensor Ks, // [..., C, 3, 3] + const CameraModelType camera_model, + // uncented transform + const UnscentedTransformParameters ut_params, + ShutterType rs_type, + const at::optional + radial_coeffs, // [..., C, 6] or [..., C, 4] optional + const at::optional tangential_coeffs, // [..., C, 2] optional + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs, // shared parameters for all cameras + // intersections + const at::Tensor tile_offsets, // [..., C, tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +); + +std::tuple +rasterize_to_pixels_from_world_3dgs_bwd( + // Gaussian parameters + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor colors, // [..., C, N, 3] or [nnz, 3] + const at::Tensor opacities, // [..., C, N] or [nnz] + const at::optional backgrounds, // [..., C, 3] + const at::optional masks, // [..., C, tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // camera + const at::Tensor viewmats0, // [..., C, 4, 4] + const at::optional + viewmats1, // [..., C, 4, 4] optional for rolling shutter + const at::Tensor Ks, // [..., C, 3, 3] + const CameraModelType camera_model, + // uncented transform + const UnscentedTransformParameters ut_params, + ShutterType rs_type, + const at::optional + radial_coeffs, // [..., C, 6] or [..., C, 4] optional + const at::optional tangential_coeffs, // [..., C, 2] optional + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs, // shared parameters for all cameras + // intersections + const at::Tensor tile_offsets, // [..., C, tile_height, tile_width] + const at::Tensor flatten_ids, // [n_isects] + // forward outputs + const at::Tensor render_alphas, // [..., C, image_height, image_width, 1] + const at::Tensor last_ids, // [..., C, image_height, image_width] + // gradients of outputs + const at::Tensor v_render_colors, // [..., C, image_height, image_width, 3] + const at::Tensor v_render_alphas // [..., C, image_height, image_width, 1] +); + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/Sycl_utils.hpp b/gsplat/sycl/include/Sycl_utils.hpp new file mode 100644 index 00000000..d312e797 --- /dev/null +++ b/gsplat/sycl/include/Sycl_utils.hpp @@ -0,0 +1,70 @@ +#pragma once + +#include + +template struct BufferType { + using type = sycl::marray; + constexpr static bool isVec{false}; +}; + +template struct BufferType { + using type = sycl::vec; + constexpr static bool isVec{true}; +}; + +template struct BufferType { + using type = sycl::vec; + constexpr static bool isVec{true}; +}; + +template struct BufferType { + using type = sycl::vec; + constexpr static bool isVec{true}; +}; + +template struct BufferType { + using type = sycl::vec; + constexpr static bool isVec{true}; +}; + +template struct BufferType { + using type = sycl::vec; + constexpr static bool isVec{true}; +}; + +template +using BufferType_t = typename BufferType::type; + +template void readToBuffer(T &dest, const void *source) { + dest = *(reinterpret_cast(source)); +} + +template void gpuAtomicAdd(T *ptr, T value) { + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device, + sycl::access::address_space::global_space> + protected_ref(*ptr); + protected_ref.fetch_add(value); +} + +template void gpuAtomicAddGlobal(T &ref, const T &value) { + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device, + sycl::access::address_space::global_space> + protected_ref(ref); + protected_ref.fetch_add(value); +} + +template void gpuAtomicAddLocal(T &ref, const T &value) { + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device, + sycl::access::address_space::local_space> + protected_ref(ref); + protected_ref.fetch_add(value); +} \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp new file mode 100644 index 00000000..32f9650e --- /dev/null +++ b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp @@ -0,0 +1,64 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "spherical_harmonics.hpp" +#include "types.hpp" + +namespace gsplat::xpu { + +template struct ComputeShBwdKernel { + const uint32_t m_N; + const uint32_t m_K; + const uint32_t m_degrees_to_use; + const vec3 *m_dirs; // [N, 3] + const T *m_coeffs; // [N, K, 3] + const bool *m_masks; // [N] + const T *m_v_colors; // [N, 3 + T *m_v_coeffs; // [N, K, 3] + T *m_v_dirs; // [N, 3] optional + + ComputeShBwdKernel( + const uint32_t N, + const uint32_t K, + const uint32_t degrees_to_use, + const vec3 *dirs, + const T *coeffs, + const bool *masks, + const T *v_colors, + T *v_coeffs, + T *v_dirs + ) + : m_N(N), m_K(K), m_degrees_to_use(degrees_to_use), m_dirs(dirs), + m_coeffs(coeffs), m_masks(masks), m_v_colors(v_colors), + m_v_coeffs(v_coeffs), m_v_dirs(v_dirs) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + if (idx >= m_N * 3) { + return; + } + uint32_t elem_id = idx / 3; + uint32_t c = idx % 3; // color channel + if (m_masks != nullptr && !m_masks[elem_id]) { + return; + } + vec3 v_dir = {0.f, 0.f, 0.f}; + sh_coeffs_to_color_fast_vjp( + m_degrees_to_use, + c, + m_dirs[elem_id], + m_coeffs + elem_id * m_K * 3, + m_v_colors + elem_id * 3, + m_v_coeffs + elem_id * m_K * 3, + m_v_dirs == nullptr ? nullptr : &v_dir + ); + + if (m_v_dirs != nullptr) { + gpuAtomicAdd(m_v_dirs + elem_id * 3, v_dir.x); + gpuAtomicAdd(m_v_dirs + elem_id * 3 + 1, v_dir.y); + gpuAtomicAdd(m_v_dirs + elem_id * 3 + 2, v_dir.z); + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp b/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp new file mode 100644 index 00000000..1abd3402 --- /dev/null +++ b/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include "spherical_harmonics.hpp" + +namespace gsplat::xpu { + +template struct ComputeShFwdKernel { + const uint32_t m_N; + const uint32_t m_K; + const uint32_t m_degrees_to_use; + const vec3 *m_dirs; // [N, 3] + const T *m_coeffs; // [N, K, 3] + const bool *m_masks; // [N] + T *m_colors; // [N, 3] + + ComputeShFwdKernel( + const uint32_t N, + const uint32_t K, + const uint32_t degrees_to_use, + const vec3 *dirs, + const T *coeffs, + const bool *masks, + T *colors + ) + : m_N(N), m_K(K), m_degrees_to_use(degrees_to_use), m_dirs(dirs), + m_coeffs(coeffs), m_masks(masks), m_colors(colors) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + if (idx >= m_N * 3) { + return; + } + uint32_t elem_id = idx / 3; + uint32_t c = idx % 3; // color channel + if (m_masks != nullptr && !m_masks[elem_id]) { + return; + } + sh_coeffs_to_color_fast( + m_degrees_to_use, + c, + m_dirs[elem_id], + m_coeffs + elem_id * m_K * 3, + m_colors + elem_id * 3 + ); + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp new file mode 100644 index 00000000..4789bc25 --- /dev/null +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp @@ -0,0 +1,289 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "proj.hpp" +#include "quat.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { + +template struct FullyFusedProjectionBwdKernel { + // fwd inputs + // New: Added B + const uint32_t m_B; + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [B, N, 3] + const T *m_covars; // [B, N, 6] optional + const T *m_quats; // [B, N, 4] optional + const T *m_scales; // [B, N, 3] optional + const T *m_viewmats; // [B, C, 4, 4] + const T *m_Ks; // [B, C, 3, 3] + const int32_t m_image_width; + const int32_t m_image_height; + const T m_eps2d; + const CameraModelType m_camera_model; + // fwd outputs + // Changed: radii is now [B, C, N, 2] + const int32_t *m_radii; // [B, C, N, 2] + const T *m_conics; // [B, C, N, 3] + const T *m_compensations; // [B, C, N] optional + // grad outputs + const T *m_v_means2d; // [B, C, N, 2] + const T *m_v_depths; // [B, C, N] + const T *m_v_conics; // [B, C, N, 3] + const T *m_v_compensations; // [B, C, N] optional + // grad inputs + T *m_v_means; // [B, N, 3] + T *m_v_covars; // [B, N, 6] optional + T *m_v_quats; // [B, N, 4] optional + T *m_v_scales; // [B, N, 3] optional + T *m_v_viewmats; // [B, C, 4, 4] optional + + FullyFusedProjectionBwdKernel( + // New: Added B + const uint32_t B, + const uint32_t C, + const uint32_t N, + const T *means, + const T *covars, + const T *quats, + const T *scales, + const T *viewmats, + const T *Ks, + const int32_t image_width, + const int32_t image_height, + const T eps2d, + const CameraModelType camera_model, + const int32_t *radii, + const T *conics, + const T *compensations, + const T *v_means2d, + const T *v_depths, + const T *v_conics, + const T *v_compensations, + T *v_means, + T *v_covars, + T *v_quats, + T *v_scales, + T *v_viewmats + ) + // New: Added m_B + : m_B(B), m_C(C), m_N(N), m_means(means), m_covars(covars), + m_quats(quats), m_scales(scales), m_viewmats(viewmats), m_Ks(Ks), + m_image_width(image_width), m_image_height(image_height), + m_eps2d(eps2d), m_camera_model(camera_model), m_radii(radii), + m_conics(conics), m_compensations(compensations), + m_v_means2d(v_means2d), m_v_depths(v_depths), m_v_conics(v_conics), + m_v_compensations(v_compensations), m_v_means(v_means), + m_v_covars(v_covars), m_v_quats(v_quats), m_v_scales(v_scales), + m_v_viewmats(v_viewmats) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + // Changed: Updated check to include B and both radii components + if (idx >= m_B * m_C * m_N || + (m_radii[idx * 2] <= 0 || m_radii[idx * 2 + 1] <= 0)) { + return; + } + + // Changed: Added bid and updated cid, gid calculation + const uint32_t bid = idx / (m_C * m_N); // batch id + const uint32_t cid = (idx / m_N) % m_C; // camera id + const uint32_t gid = idx % m_N; // gaussian id + + // Changed: Updated pointer arithmetic to include B + const T *means = m_means + bid * m_N * 3 + gid * 3; + const T *viewmats = m_viewmats + bid * m_C * 16 + cid * 16; + const T *Ks = m_Ks + bid * m_C * 9 + cid * 9; + const T *conics = m_conics + idx * 3; + const T *v_means2d = m_v_means2d + idx * 2; + const T *v_depths = m_v_depths + idx; + const T *v_conics = m_v_conics + idx * 3; + + // vjp: compute the inverse of the 2d covariance + mat2 covar2d_inv = + mat2(conics[0], conics[1], conics[1], conics[2]); + mat2 v_covar2d_inv = mat2( + v_conics[0], v_conics[1] * .5f, v_conics[1] * .5f, v_conics[2] + ); + mat2 v_covar2d(0.f); + inverse_vjp(covar2d_inv, v_covar2d_inv, v_covar2d); + + if (m_v_compensations != nullptr) { + // vjp: compensation term + const T compensation = m_compensations[idx]; + const T v_compensation = m_v_compensations[idx]; + add_blur_vjp( + m_eps2d, covar2d_inv, compensation, v_compensation, v_covar2d + ); + } + + // transform Gaussian to camera space + mat3 R = mat3( + viewmats[0], + viewmats[4], + viewmats[8], // 1st column + viewmats[1], + viewmats[5], + viewmats[9], // 2nd column + viewmats[2], + viewmats[6], + viewmats[10] // 3rd column + ); + vec3 t = vec3(viewmats[3], viewmats[7], viewmats[11]); + + mat3 covar; + vec4 quat; + vec3 scale; + if (m_covars != nullptr) { + // Changed: Updated pointer arithmetic + const T *covars = m_covars + bid * m_N * 6 + gid * 6; + covar = mat3( + covars[0], + covars[1], + covars[2], // 1st column + covars[1], + covars[3], + covars[4], // 2nd column + covars[2], + covars[4], + covars[5] // 3rd column + ); + } else { + // compute from quaternions and scales + // Changed: Updated pointer arithmetic + quat = glm::make_vec4(m_quats + bid * m_N * 4 + gid * 4); + scale = glm::make_vec3(m_scales + bid * m_N * 3 + gid * 3); + quat_scale_to_covar_preci(quat, scale, &covar, nullptr); + } + vec3 mean_c; + pos_world_to_cam(R, t, glm::make_vec3(means), mean_c); + mat3 covar_c; + covar_world_to_cam(R, covar, covar_c); + + // vjp: perspective projection + T fx = Ks[0], cx = Ks[2], fy = Ks[4], cy = Ks[5]; + mat3 v_covar_c(0.f); + vec3 v_mean_c(0.f); + + switch (m_camera_model) { + case CameraModelType::PINHOLE: // perspective projection + persp_proj_vjp( + mean_c, + covar_c, + fx, + fy, + cx, + cy, + m_image_width, + m_image_height, + v_covar2d, + glm::make_vec2(v_means2d), + v_mean_c, + v_covar_c + ); + break; + case CameraModelType::ORTHO: // orthographic projection + ortho_proj_vjp( + mean_c, + covar_c, + fx, + fy, + cx, + cy, + m_image_width, + m_image_height, + v_covar2d, + glm::make_vec2(v_means2d), + v_mean_c, + v_covar_c + ); + break; + case CameraModelType::FISHEYE: // fisheye projection + fisheye_proj_vjp( + mean_c, + covar_c, + fx, + fy, + cx, + cy, + m_image_width, + m_image_height, + v_covar2d, + glm::make_vec2(v_means2d), + v_mean_c, + v_covar_c + ); + break; + } + + // add contribution from v_depths + v_mean_c.z += v_depths[0]; + + // vjp: transform Gaussian covariance to camera space + vec3 v_mean(0.f); + mat3 v_covar(0.f); + mat3 v_R(0.f); + vec3 v_t(0.f); + pos_world_to_cam_vjp( + R, t, glm::make_vec3(means), v_mean_c, v_R, v_t, v_mean + ); + covar_world_to_cam_vjp(R, covar, v_covar_c, v_R, v_covar); + + if (m_v_means != nullptr) { + // Changed: Updated pointer arithmetic + T *v_means = m_v_means + bid * m_N * 3 + gid * 3; +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { + gpuAtomicAdd(v_means + i, v_mean[i]); + } + } + + if (m_v_covars != nullptr) { + // Changed: Updated pointer arithmetic + T *v_covars = m_v_covars + bid * m_N * 6 + gid * 6; + gpuAtomicAdd(v_covars, v_covar[0][0]); + gpuAtomicAdd(v_covars + 1, v_covar[0][1] + v_covar[1][0]); + gpuAtomicAdd(v_covars + 2, v_covar[0][2] + v_covar[2][0]); + gpuAtomicAdd(v_covars + 3, v_covar[1][1]); + gpuAtomicAdd(v_covars + 4, v_covar[1][2] + v_covar[2][1]); + gpuAtomicAdd(v_covars + 5, v_covar[2][2]); + } else { + // Directly output gradients w.r.t. the quaternion and scale + mat3 rotmat = quat_to_rotmat(quat); + vec4 v_quat(0.f); + vec3 v_scale(0.f); + quat_scale_to_covar_vjp( + quat, scale, rotmat, v_covar, v_quat, v_scale + ); + // Changed: Updated pointer arithmetic + T *v_quats = m_v_quats + bid * m_N * 4 + gid * 4; + T *v_scales = m_v_scales + bid * m_N * 3 + gid * 3; + gpuAtomicAdd(v_quats, v_quat[0]); + gpuAtomicAdd(v_quats + 1, v_quat[1]); + gpuAtomicAdd(v_quats + 2, v_quat[2]); + gpuAtomicAdd(v_quats + 3, v_quat[3]); + gpuAtomicAdd(v_scales, v_scale[0]); + gpuAtomicAdd(v_scales + 1, v_scale[1]); + gpuAtomicAdd(v_scales + 2, v_scale[2]); + } + + if (m_v_viewmats != nullptr) { + // Changed: Updated pointer arithmetic + T *v_viewmats = m_v_viewmats + bid * m_C * 16 + cid * 16; +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + gpuAtomicAdd(v_viewmats + i * 4 + j, v_R[j][i]); + } + gpuAtomicAdd(v_viewmats + i * 4 + 3, v_t[i]); + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp new file mode 100644 index 00000000..951d3b5e --- /dev/null +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp @@ -0,0 +1,255 @@ +#pragma once + +#include "proj.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { + +template struct FullyFusedProjectionFwdKernel { + // New: Added B + const uint32_t m_B; + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [B, N, 3] + const T *m_covars; // [B, N, 6] optional + const T *m_quats; // [B, N, 4] optional + const T *m_scales; // [B, N, 3] optional + // New: Added opacities + const T *m_opacities; // [B, N] optional + const T *m_viewmats; // [B, C, 4, 4] + const T *m_Ks; // [B, C, 3, 3] + const int32_t m_image_width; + const int32_t m_image_height; + const T m_eps2d; + const T m_near_plane; + const T m_far_plane; + const T m_radius_clip; + const CameraModelType m_camera_model; + // outputs + // Changed: radii is now [B, C, N, 2] + int32_t *m_radii; // [B, C, N, 2] + T *m_means2d; // [B, C, N, 2] + T *m_depths; // [B, C, N] + T *m_conics; // [B, C, N, 3] + T *m_compensations; // [B, C, N] optional + + FullyFusedProjectionFwdKernel( + // New: Added B + const uint32_t B, + const uint32_t C, + const uint32_t N, + const T *means, + const T *covars, + const T *quats, + const T *scales, + // New: Added opacities + const T *opacities, + const T *viewmats, + const T *Ks, + const int32_t image_width, + const int32_t image_height, + const T eps2d, + const T near_plane, + const T far_plane, + const T radius_clip, + const CameraModelType camera_model, + int32_t *radii, + T *means2d, + T *depths, + T *conics, + T *compensations + ) + // New: Added m_B and m_opacities + : m_B(B), m_C(C), m_N(N), m_means(means), m_covars(covars), + m_quats(quats), m_scales(scales), m_opacities(opacities), + m_viewmats(viewmats), m_Ks(Ks), m_image_width(image_width), + m_image_height(image_height), m_eps2d(eps2d), + m_near_plane(near_plane), m_far_plane(far_plane), + m_radius_clip(radius_clip), m_camera_model(camera_model), + m_radii(radii), m_means2d(means2d), m_depths(depths), + m_conics(conics), m_compensations(compensations) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + // Changed: Updated upper bound to include B + if (idx >= m_B * m_C * m_N) { + return; + } + // Changed: Added bid and updated cid, gid calculation + const uint32_t bid = idx / (m_C * m_N); // batch id + const uint32_t cid = (idx / m_N) % m_C; // camera id + const uint32_t gid = idx % m_N; // gaussian id + + // Changed: Updated pointer arithmetic to include B + const T *means = m_means + bid * m_N * 3 + gid * 3; + const T *viewmats = m_viewmats + bid * m_C * 16 + cid * 16; + const T *Ks = m_Ks + bid * m_C * 9 + cid * 9; + + // glm is column-major but input is row-major + mat3 R = mat3( + viewmats[0], + viewmats[4], + viewmats[8], // 1st column + viewmats[1], + viewmats[5], + viewmats[9], // 2nd column + viewmats[2], + viewmats[6], + viewmats[10] // 3rd column + ); + vec3 t = vec3(viewmats[3], viewmats[7], viewmats[11]); + + // transform Gaussian center to camera space + vec3 mean_c; + pos_world_to_cam(R, t, glm::make_vec3(means), mean_c); + if (mean_c.z < m_near_plane || mean_c.z > m_far_plane) { + // Changed: Set both radii to 0 + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + // transform Gaussian covariance to camera space + mat3 covar; + if (m_covars != nullptr) { + // Changed: Updated pointer arithmetic + const T *covars = m_covars + bid * m_N * 6 + gid * 6; + covar = mat3( + covars[0], + covars[1], + covars[2], // 1st column + covars[1], + covars[3], + covars[4], // 2nd column + covars[2], + covars[4], + covars[5] // 3rd column + ); + } else { + // compute from quaternions and scales + // Changed: Updated pointer arithmetic + const T *quats = m_quats + bid * m_N * 4 + gid * 4; + const T *scales = m_scales + bid * m_N * 3 + gid * 3; + quat_scale_to_covar_preci( + glm::make_vec4(quats), glm::make_vec3(scales), &covar, nullptr + ); + } + mat3 covar_c; + covar_world_to_cam(R, covar, covar_c); + + // perspective projection + mat2 covar2d; + vec2 mean2d; + + switch (m_camera_model) { + case CameraModelType::PINHOLE: // perspective projection + persp_proj( + mean_c, + covar_c, + Ks[0], + Ks[4], + Ks[2], + Ks[5], + m_image_width, + m_image_height, + covar2d, + mean2d + ); + break; + case CameraModelType::ORTHO: // orthographic projection + ortho_proj( + mean_c, + covar_c, + Ks[0], + Ks[4], + Ks[2], + Ks[5], + m_image_width, + m_image_height, + covar2d, + mean2d + ); + break; + case CameraModelType::FISHEYE: // fisheye projection + fisheye_proj( + mean_c, + covar_c, + Ks[0], + Ks[4], + Ks[2], + Ks[5], + m_image_width, + m_image_height, + covar2d, + mean2d + ); + break; + } + + T compensation; + T det = add_blur(m_eps2d, covar2d, compensation); + if (det <= 0.f) { + // Changed: Set both radii to 0 + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + // compute the inverse of the 2d covariance + mat2 covar2d_inv; + inverse(covar2d, covar2d_inv); + + // New: Opacity-aware bounding box and radius calculation + const T ALPHA_THRESHOLD = 1.f / 255.f; + T extend = 3.33f; + if (m_opacities != nullptr) { + T opacity = m_opacities[bid * m_N + gid]; + if (m_compensations != nullptr) { + opacity *= compensation; + } + if (opacity < ALPHA_THRESHOLD) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + extend = sycl::min( + extend, sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD)) + ); + } + + T radius_x = sycl::ceil(extend * sycl::sqrt(covar2d[0][0])); + T radius_y = sycl::ceil(extend * sycl::sqrt(covar2d[1][1])); + + if (radius_x <= m_radius_clip && radius_y <= m_radius_clip) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + // mask out gaussians outside the image region + if (mean2d.x + radius_x <= 0 || mean2d.x - radius_x >= m_image_width || + mean2d.y + radius_y <= 0 || mean2d.y - radius_y >= m_image_height) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + // write to outputs + // Changed: Write radius_x and radius_y + m_radii[idx * 2] = (int32_t)radius_x; + m_radii[idx * 2 + 1] = (int32_t)radius_y; + m_means2d[idx * 2] = mean2d.x; + m_means2d[idx * 2 + 1] = mean2d.y; + m_depths[idx] = mean_c.z; + m_conics[idx * 3] = covar2d_inv[0][0]; + m_conics[idx * 3 + 1] = covar2d_inv[0][1]; + m_conics[idx * 3 + 2] = covar2d_inv[1][1]; + if (m_compensations != nullptr) { + m_compensations[idx] = compensation; + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp b/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp new file mode 100644 index 00000000..92f968b0 --- /dev/null +++ b/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp @@ -0,0 +1,65 @@ +#pragma once + +namespace gsplat::xpu { + +struct IsectOffsetEncodeKernel { + + const uint32_t m_n_isects; + const int64_t *m_isect_ids; + const uint32_t m_C; + const uint32_t m_n_tiles; + const uint32_t m_tile_n_bits; + int32_t *m_offsets; //[C, n_tiles] + + IsectOffsetEncodeKernel( + const uint32_t n_isects, + const int64_t *isect_ids, + const uint32_t C, + const uint32_t n_tiles, + const uint32_t tile_n_bits, + int32_t *offsets + ) + : m_n_isects(n_isects), m_isect_ids(isect_ids), m_C(C), + m_n_tiles(n_tiles), m_tile_n_bits(tile_n_bits), m_offsets(offsets) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + + if (idx >= m_n_isects) + return; + + int64_t isect_id_curr = m_isect_ids[idx] >> 32; + int64_t cid_curr = isect_id_curr >> m_tile_n_bits; + int64_t tid_curr = isect_id_curr & ((1 << m_tile_n_bits) - 1); + int64_t id_curr = cid_curr * m_n_tiles + tid_curr; + + if (idx == 0) { + // write out the offsets until the first valid tile (inclusive) + for (uint32_t i = 0; i < id_curr + 1; ++i) + m_offsets[i] = static_cast(idx); + } + if (idx == m_n_isects - 1) { + // write out the rest of the offsets + for (uint32_t i = id_curr + 1; i < m_C * m_n_tiles; ++i) + m_offsets[i] = static_cast(m_n_isects); + } + + if (idx > 0) { + // visit the current and previous isect_id and check if the (cid, + // tile_id) pair changes. + int64_t isect_id_prev = + m_isect_ids[idx - 1] >> 32; // shift out the depth + if (isect_id_prev == isect_id_curr) + return; + + // write out the offsets between the previous and current tiles + int64_t cid_prev = isect_id_prev >> m_tile_n_bits; + int64_t tid_prev = isect_id_prev & ((1 << m_tile_n_bits) - 1); + int64_t id_prev = cid_prev * m_n_tiles + tid_prev; + for (uint32_t i = id_prev + 1; i < id_curr + 1; ++i) + m_offsets[i] = static_cast(idx); + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/IsectTilesKernel.hpp b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp new file mode 100644 index 00000000..d5cb93a7 --- /dev/null +++ b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp @@ -0,0 +1,151 @@ +#pragma once + +#include "transform.hpp" +#include "types.hpp" +#include "utils.hpp" +#include + +namespace gsplat::xpu { + +struct uint2 { + uint32_t x; + uint32_t y; +}; + +template struct IsectTilesKernel { + const bool m_packed; + const uint32_t m_C; + const uint32_t m_N; + const uint32_t m_nnz; + const int64_t *m_camera_ids; // [nnz] optional + const int64_t *m_gaussian_ids; // [nnz] optional + const T *m_means2d; // [C, N, 2] or [nnz, 2] + const int32_t *m_radii; // [C, N] or [nnz] + const T *m_depths; // [C, N] or [nnz] + const int64_t *m_cum_tiles_per_gauss; // [C, N] or [nnz] + const uint32_t m_tile_size; + const uint32_t m_tile_width; + const uint32_t m_tile_height; + const uint32_t m_tile_n_bits; + int32_t *m_tiles_per_gauss; // [C, N] or [nnz] + int64_t *m_isect_ids; // [n_isects] + int32_t *m_flatten_ids; // [n_isects] + + IsectTilesKernel( + const bool packed, + const uint32_t C, + const uint32_t N, + const uint32_t nnz, + const int64_t *camera_ids, + const int64_t *gaussian_ids, + const T *means2d, + const int32_t *radii, + const T *depths, + const int64_t *cum_tiles_per_gauss, + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + const uint32_t tile_n_bits, + int32_t *tiles_per_gauss, + int64_t *isect_ids, + int32_t *flatten_ids + ) + : m_packed(packed), m_C(C), m_N(N), m_nnz(nnz), + m_camera_ids(camera_ids), m_gaussian_ids(gaussian_ids), + m_means2d(means2d), m_radii(radii), m_depths(depths), + m_cum_tiles_per_gauss(cum_tiles_per_gauss), m_tile_size(tile_size), + m_tile_width(tile_width), m_tile_height(tile_height), + m_tile_n_bits(tile_n_bits), m_tiles_per_gauss(tiles_per_gauss), + m_isect_ids(isect_ids), m_flatten_ids(flatten_ids) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + + bool first_pass = m_cum_tiles_per_gauss == nullptr; + if (idx >= (m_packed ? m_nnz : m_C * m_N)) { + return; + } + + const T radius_x = m_radii[idx * 2]; + const T radius_y = m_radii[idx * 2 + 1]; + if (radius_x <= 0 || radius_y <= 0) { + if (first_pass) { + m_tiles_per_gauss[idx] = 0; + } + return; + } + + vec2 mean2d = glm::make_vec2(m_means2d + 2 * idx); + + T tile_radius_x = radius_x / static_cast(m_tile_size); + T tile_radius_y = radius_y / static_cast(m_tile_size); + T tile_x = mean2d.x / static_cast(m_tile_size); + T tile_y = mean2d.y / static_cast(m_tile_size); + + uint2 tile_min, tile_max; + // Use the separate x and y tile radii to calculate the bounding box. + tile_min.x = sycl::min( + sycl::max( + (uint32_t)0, (uint32_t)sycl::floor(tile_x - tile_radius_x) + ), + m_tile_width + ); + tile_min.y = sycl::min( + sycl::max( + (uint32_t)0, (uint32_t)sycl::floor(tile_y - tile_radius_y) + ), + m_tile_height + ); + tile_max.x = sycl::min( + sycl::max( + (uint32_t)0, (uint32_t)sycl::ceil(tile_x + tile_radius_x) + ), + m_tile_width + ); + tile_max.y = sycl::min( + sycl::max( + (uint32_t)0, (uint32_t)sycl::ceil(tile_y + tile_radius_y) + ), + m_tile_height + ); + + if (first_pass) { + // first pass only writes out tiles_per_gauss + m_tiles_per_gauss[idx] = static_cast( + (tile_max.y - tile_min.y) * (tile_max.x - tile_min.x) + ); + return; + } + + int64_t cid; // camera id + if (m_packed) { + // parallelize over nnz + cid = m_camera_ids[idx]; + // gid = gaussian_ids[idx]; + } else { + // parallelize over C * N + cid = idx / m_N; + // gid = idx % N; + } + + const int64_t cid_enc = cid << (32 + m_tile_n_bits); + + int32_t depth_i32 = *reinterpret_cast(&m_depths[idx]); + int64_t depth_id_enc = static_cast(depth_i32); + + int64_t cur_idx = (idx == 0) ? 0 : m_cum_tiles_per_gauss[idx - 1]; + for (int32_t i = tile_min.y; i < tile_max.y; ++i) { + for (int32_t j = tile_min.x; j < tile_max.x; ++j) { + int64_t tile_id = i * m_tile_width + j; + // e.g. tile_n_bits = 22: + // camera id (10 bits) | tile id (22 bits) | depth (32 bits) + m_isect_ids[cur_idx] = cid_enc | (tile_id << 32) | depth_id_enc; + // the flatten index in [C * N] or [nnz] + m_flatten_ids[cur_idx] = static_cast(idx); + ++cur_idx; + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp new file mode 100644 index 00000000..0b8b96bd --- /dev/null +++ b/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp @@ -0,0 +1,316 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "proj.hpp" +#include "quat.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" +#include "utils.hpp" +#include + +namespace gsplat::xpu { + +template struct PackedProjectionBwdKernel { + // fwd inputs + const uint32_t m_B; + const uint32_t m_C; + const uint32_t m_N; + const uint32_t m_nnz; + const T *m_means; + const T *m_covars; + const T *m_quats; + const T *m_scales; + const T *m_viewmats; + const T *m_Ks; + const int32_t m_image_width; + const int32_t m_image_height; + const T m_eps2d; + const CameraModelType m_camera_model; + // fwd outputs (packed) + const int64_t *m_batch_ids; + const int64_t *m_camera_ids; + const int64_t *m_gaussian_ids; + const T *m_conics; + const T *m_compensations; + // grad outputs (packed) + const T *m_v_means2d; + const T *m_v_depths; + const T *m_v_conics; + const T *m_v_compensations; + const bool m_sparse_grad; + // grad inputs + T *m_v_means; + T *m_v_covars; + T *m_v_quats; + T *m_v_scales; + T *m_v_viewmats; + + PackedProjectionBwdKernel( + uint32_t B, + uint32_t C, + uint32_t N, + uint32_t nnz, + const T *means, + const T *covars, + const T *quats, + const T *scales, + const T *viewmats, + const T *Ks, + int32_t image_width, + int32_t image_height, + T eps2d, + CameraModelType camera_model, + const int64_t *batch_ids, + const int64_t *camera_ids, + const int64_t *gaussian_ids, + const T *conics, + const T *compensations, + const T *v_means2d, + const T *v_depths, + const T *v_conics, + const T *v_compensations, + bool sparse_grad, + T *v_means, + T *v_covars, + T *v_quats, + T *v_scales, + T *v_viewmats + ) + : m_B(B), m_C(C), m_N(N), m_nnz(nnz), m_means(means), m_covars(covars), + m_quats(quats), m_scales(scales), m_viewmats(viewmats), m_Ks(Ks), + m_image_width(image_width), m_image_height(image_height), + m_eps2d(eps2d), m_camera_model(camera_model), m_batch_ids(batch_ids), + m_camera_ids(camera_ids), m_gaussian_ids(gaussian_ids), + m_conics(conics), m_compensations(compensations), + m_v_means2d(v_means2d), m_v_depths(v_depths), m_v_conics(v_conics), + m_v_compensations(v_compensations), m_sparse_grad(sparse_grad), + m_v_means(v_means), m_v_covars(v_covars), m_v_quats(v_quats), + m_v_scales(v_scales), m_v_viewmats(v_viewmats) {} + + void operator()(sycl::nd_item<1> item) const { + uint32_t idx = item.get_global_id(0); + if (idx >= m_nnz) { + return; + } + + const int64_t bid = m_batch_ids[idx]; + const int64_t cid = m_camera_ids[idx]; + const int64_t gid = m_gaussian_ids[idx]; + + // --- VJP Calculation (same as fused, but with packed inputs) --- + + mat2 v_covar2d(0.f); + { + const T *conics = m_conics + idx * 3; + const T *v_conics = m_v_conics + idx * 3; + mat2 covar2d_inv = + mat2(conics[0], conics[1], conics[1], conics[2]); + mat2 v_covar2d_inv = mat2( + v_conics[0], v_conics[1] * 0.5f, v_conics[1] * 0.5f, v_conics[2] + ); + inverse_vjp(covar2d_inv, v_covar2d_inv, v_covar2d); + + if (m_v_compensations != nullptr) { + const T compensation = m_compensations[idx]; + const T v_compensation = m_v_compensations[idx]; + add_blur_vjp( + m_eps2d, + covar2d_inv, + compensation, + v_compensation, + v_covar2d + ); + } + } + + const T *means = m_means + bid * m_N * 3 + gid * 3; + const T *viewmats = m_viewmats + bid * m_C * 16 + cid * 16; + const T *Ks = m_Ks + bid * m_C * 9 + cid * 9; + + mat3 R( + viewmats[0], + viewmats[4], + viewmats[8], + viewmats[1], + viewmats[5], + viewmats[9], + viewmats[2], + viewmats[6], + viewmats[10] + ); + vec3 t(viewmats[3], viewmats[7], viewmats[11]); + + mat3 covar; + vec4 quat; + vec3 scale; + if (m_covars != nullptr) { + const T *covars = m_covars + bid * m_N * 6 + gid * 6; + covar = mat3( + covars[0], + covars[1], + covars[2], + covars[1], + covars[3], + covars[4], + covars[2], + covars[4], + covars[5] + ); + } else { + quat = glm::make_vec4(m_quats + bid * m_N * 4 + gid * 4); + scale = glm::make_vec3(m_scales + bid * m_N * 3 + gid * 3); + quat_scale_to_covar_preci(quat, scale, &covar, nullptr); + } + + vec3 mean_c; + pos_world_to_cam(R, t, glm::make_vec3(means), mean_c); + mat3 covar_c; + covar_world_to_cam(R, covar, covar_c); + + mat3 v_covar_c(0.f); + vec3 v_mean_c(0.f); + const T *v_means2d = m_v_means2d + idx * 2; + + switch (m_camera_model) { + case CameraModelType::PINHOLE: + persp_proj_vjp( + mean_c, + covar_c, + Ks[0], + Ks[4], + Ks[2], + Ks[5], + m_image_width, + m_image_height, + v_covar2d, + glm::make_vec2(v_means2d), + v_mean_c, + v_covar_c + ); + break; + case CameraModelType::ORTHO: + ortho_proj_vjp( + mean_c, + covar_c, + Ks[0], + Ks[4], + Ks[2], + Ks[5], + m_image_width, + m_image_height, + v_covar2d, + glm::make_vec2(v_means2d), + v_mean_c, + v_covar_c + ); + break; + case CameraModelType::FISHEYE: + fisheye_proj_vjp( + mean_c, + covar_c, + Ks[0], + Ks[4], + Ks[2], + Ks[5], + m_image_width, + m_image_height, + v_covar2d, + glm::make_vec2(v_means2d), + v_mean_c, + v_covar_c + ); + break; + } + + v_mean_c.z += m_v_depths[idx]; + + vec3 v_mean(0.f); + mat3 v_covar(0.f); + mat3 v_R(0.f); + vec3 v_t(0.f); + pos_world_to_cam_vjp( + R, t, glm::make_vec3(means), v_mean_c, v_R, v_t, v_mean + ); + covar_world_to_cam_vjp(R, covar, v_covar_c, v_R, v_covar); + + // --- Gradient Accumulation --- + + if (m_sparse_grad) { + // Write gradients to sparse output tensors (no atomics needed) + if (m_v_means != nullptr) { + T *v_means_out = m_v_means + idx * 3; + v_means_out[0] = v_mean.x; + v_means_out[1] = v_mean.y; + v_means_out[2] = v_mean.z; + } + if (m_v_covars != nullptr) { + T *v_covars_out = m_v_covars + idx * 6; + v_covars_out[0] = v_covar[0][0]; + v_covars_out[1] = v_covar[0][1] + v_covar[1][0]; + v_covars_out[2] = v_covar[0][2] + v_covar[2][0]; + v_covars_out[3] = v_covar[1][1]; + v_covars_out[4] = v_covar[1][2] + v_covar[2][1]; + v_covars_out[5] = v_covar[2][2]; + } else { + mat3 rotmat = quat_to_rotmat(quat); + vec4 v_quat(0.f); + vec3 v_scale(0.f); + quat_scale_to_covar_vjp( + quat, scale, rotmat, v_covar, v_quat, v_scale + ); + T *v_quats_out = m_v_quats + idx * 4; + T *v_scales_out = m_v_scales + idx * 3; + v_quats_out[0] = v_quat.x; + v_quats_out[1] = v_quat.y; + v_quats_out[2] = v_quat.z; + v_quats_out[3] = v_quat.w; + v_scales_out[0] = v_scale.x; + v_scales_out[1] = v_scale.y; + v_scales_out[2] = v_scale.z; + } + } else { + // Atomically accumulate gradients into dense tensors + if (m_v_means != nullptr) { + T *v_means_out = m_v_means + bid * m_N * 3 + gid * 3; + for (int i = 0; i < 3; ++i) { + gpuAtomicAdd(&v_means_out[i], v_mean[i]); + } + } + if (m_v_covars != nullptr) { + T *v_covars_out = m_v_covars + bid * m_N * 6 + gid * 6; + gpuAtomicAdd(&v_covars_out[0], v_covar[0][0]); + gpuAtomicAdd(&v_covars_out[1], v_covar[0][1] + v_covar[1][0]); + gpuAtomicAdd(&v_covars_out[2], v_covar[0][2] + v_covar[2][0]); + gpuAtomicAdd(&v_covars_out[3], v_covar[1][1]); + gpuAtomicAdd(&v_covars_out[4], v_covar[1][2] + v_covar[2][1]); + gpuAtomicAdd(&v_covars_out[5], v_covar[2][2]); + } else { + mat3 rotmat = quat_to_rotmat(quat); + vec4 v_quat(0.f); + vec3 v_scale(0.f); + quat_scale_to_covar_vjp( + quat, scale, rotmat, v_covar, v_quat, v_scale + ); + T *v_quats_out = m_v_quats + bid * m_N * 4 + gid * 4; + T *v_scales_out = m_v_scales + bid * m_N * 3 + gid * 3; + for (int i = 0; i < 4; ++i) + gpuAtomicAdd(&v_quats_out[i], v_quat[i]); + for (int i = 0; i < 3; ++i) + gpuAtomicAdd(&v_scales_out[i], v_scale[i]); + } + } + + // v_viewmats is always dense and requires atomics + if (m_v_viewmats != nullptr) { + T *v_viewmats_out = m_v_viewmats + bid * m_C * 16 + cid * 16; + for (uint32_t i = 0; i < 3; i++) { // rows + for (uint32_t j = 0; j < 3; j++) { // cols + gpuAtomicAdd(&v_viewmats_out[i * 4 + j], v_R[j][i]); + } + gpuAtomicAdd(&v_viewmats_out[i * 4 + 3], v_t[i]); + } + } + } +}; + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp new file mode 100644 index 00000000..2d218747 --- /dev/null +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -0,0 +1,326 @@ +#pragma once + +#include "proj.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" +#include "utils.hpp" +#include + +namespace gsplat::xpu { + +template struct PackedProjectionFwdKernel { + // Inputs + const uint32_t m_B; + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; + const T *m_covars; + const T *m_quats; + const T *m_scales; + const T *m_opacities; + const T *m_viewmats; + const T *m_Ks; + const int32_t m_image_width; + const int32_t m_image_height; + const T m_eps2d; + const T m_near_plane; + const T m_far_plane; + const T m_radius_clip; + const CameraModelType m_camera_model; + const int64_t *m_block_accum; // Packing helper for the second pass + + // Outputs + int32_t *m_block_cnts; + int32_t *m_indptr; + int64_t *m_batch_ids; + int64_t *m_camera_ids; + int64_t *m_gaussian_ids; + int32_t *m_radii; + T *m_means2d; + T *m_depths; + T *m_conics; + T *m_compensations; + + PackedProjectionFwdKernel( + uint32_t B, + uint32_t C, + uint32_t N, + const T *means, + const T *covars, + const T *quats, + const T *scales, + const T *opacities, + const T *viewmats, + const T *Ks, + int32_t image_width, + int32_t image_height, + T eps2d, + T near_plane, + T far_plane, + T radius_clip, + CameraModelType camera_model, + const int64_t *block_accum, + // outputs + int32_t *block_cnts, + int32_t *indptr, + int64_t *batch_ids, + int64_t *camera_ids, + int64_t *gaussian_ids, + int32_t *radii, + T *means2d, + T *depths, + T *conics, + T *compensations + ) + : m_B(B), m_C(C), m_N(N), m_means(means), m_covars(covars), + m_quats(quats), m_scales(scales), m_opacities(opacities), + m_viewmats(viewmats), m_Ks(Ks), m_image_width(image_width), + m_image_height(image_height), m_eps2d(eps2d), + m_near_plane(near_plane), m_far_plane(far_plane), + m_radius_clip(radius_clip), m_camera_model(camera_model), + m_block_accum(block_accum), m_block_cnts(block_cnts), + m_indptr(indptr), m_batch_ids(batch_ids), m_camera_ids(camera_ids), + m_gaussian_ids(gaussian_ids), m_radii(radii), m_means2d(means2d), + m_depths(depths), m_conics(conics), m_compensations(compensations) {} + + void operator()(sycl::nd_item<2> item) const { + auto group = item.get_group(); + + sycl::id<2> group_id = item.get_group().get_group_id(); + sycl::range<2> group_range = item.get_group_range(); + sycl::id<2> local_id_2d = item.get_local_id(); + sycl::range<2> local_range = item.get_local_range(); + + int32_t blocks_per_row = + group_range[1]; // Get range of the 2nd dimension + + int32_t row_idx = group_id[0]; // Get group ID of the 1st dimension + int32_t block_col_idx = + group_id[1]; // Get group ID of the 2nd dimension + int32_t block_idx = row_idx * blocks_per_row + block_col_idx; + + int32_t local_id = local_id_2d[1]; // Get local ID of the 2nd dimension + int32_t col_idx = block_col_idx * local_range[1] + local_id; + + const int32_t bid = row_idx / m_C; + const int32_t cid = row_idx % m_C; + const int32_t gid = col_idx; + + bool valid = (bid < m_B) && (cid < m_C) && (gid < m_N); + + // --- Culling logic shared between both passes --- + vec3 mean_c; + mat3 R; + if (valid) { + const T *current_means = m_means + bid * m_N * 3 + gid * 3; + const T *current_viewmats = m_viewmats + bid * m_C * 16 + cid * 16; + + R = mat3( + current_viewmats[0], + current_viewmats[4], + current_viewmats[8], + current_viewmats[1], + current_viewmats[5], + current_viewmats[9], + current_viewmats[2], + current_viewmats[6], + current_viewmats[10] + ); + vec3 t( + current_viewmats[3], current_viewmats[7], current_viewmats[11] + ); + + pos_world_to_cam(R, t, glm::make_vec3(current_means), mean_c); + if (mean_c.z < m_near_plane || mean_c.z > m_far_plane) { + valid = false; + } + } + + mat2 covar2d; + vec2 mean2d; + mat2 covar2d_inv; + T compensation; + if (valid) { + mat3 covar; + if (m_covars != nullptr) { + const T *current_covars = m_covars + bid * m_N * 6 + gid * 6; + covar = mat3( + current_covars[0], + current_covars[1], + current_covars[2], + current_covars[1], + current_covars[3], + current_covars[4], + current_covars[2], + current_covars[4], + current_covars[5] + ); + } else { + const T *current_quats = m_quats + bid * m_N * 4 + gid * 4; + const T *current_scales = m_scales + bid * m_N * 3 + gid * 3; + quat_scale_to_covar_preci( + glm::make_vec4(current_quats), + glm::make_vec3(current_scales), + &covar, + nullptr + ); + } + mat3 covar_c; + covar_world_to_cam(R, covar, covar_c); + + const T *current_Ks = m_Ks + bid * m_C * 9 + cid * 9; + switch (m_camera_model) { + case CameraModelType::PINHOLE: + persp_proj( + mean_c, + covar_c, + current_Ks[0], + current_Ks[4], + current_Ks[2], + current_Ks[5], + m_image_width, + m_image_height, + covar2d, + mean2d + ); + break; + case CameraModelType::ORTHO: + ortho_proj( + mean_c, + covar_c, + current_Ks[0], + current_Ks[4], + current_Ks[2], + current_Ks[5], + m_image_width, + m_image_height, + covar2d, + mean2d + ); + break; + case CameraModelType::FISHEYE: + fisheye_proj( + mean_c, + covar_c, + current_Ks[0], + current_Ks[4], + current_Ks[2], + current_Ks[5], + m_image_width, + m_image_height, + covar2d, + mean2d + ); + break; + } + + T det = add_blur(m_eps2d, covar2d, compensation); + if (det <= 0.f) { + valid = false; + } else { + inverse(covar2d, covar2d_inv); + } + } + + T radius_x, radius_y; + if (valid) { + const T ALPHA_THRESHOLD = 1.f / 255.f; + T extend = 3.33f; + if (m_opacities != nullptr) { + T opacity = m_opacities[bid * m_N + gid]; + if (m_compensations != nullptr) { + opacity *= compensation; + } + if (opacity < ALPHA_THRESHOLD) { + valid = false; + } + extend = sycl::fmin( + extend, + sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD)) + ); + } + + radius_x = sycl::ceil(extend * sycl::sqrt(covar2d[0][0])); + radius_y = sycl::ceil(extend * sycl::sqrt(covar2d[1][1])); + + if (radius_x <= m_radius_clip && radius_y <= m_radius_clip) { + valid = false; + } + + if (mean2d.x + radius_x <= 0 || + mean2d.x - radius_x >= m_image_width || + mean2d.y + radius_y <= 0 || + mean2d.y - radius_y >= m_image_height) { + valid = false; + } + } + + // --- Pass-specific logic --- + + if (m_block_cnts != nullptr) { + // First pass: Count visible Gaussians in this block. + int32_t thread_data = static_cast(valid); + bool any_valid = sycl::any_of_group(group, valid); + if (any_valid) { + // Reduce the count of valid Gaussians across the work-group. + int32_t aggregate = + sycl::reduce_over_group(group, thread_data, sycl::plus<>()); + if (local_id == 0) { + m_block_cnts[block_idx] = aggregate; + } + } else { + if (local_id == 0) { + m_block_cnts[block_idx] = 0; + } + } + + } else { + // Second pass: Write data for visible Gaussians. + int64_t thread_data = static_cast(valid); + bool any_valid = sycl::any_of_group(group, valid); + if (any_valid) { + // Perform an exclusive scan to find the local offset for this + // thread. + int64_t local_offset = sycl::exclusive_scan_over_group( + group, thread_data, sycl::plus<>() + ); + + if (valid) { + int64_t global_offset = local_offset; + if (block_idx > 0) { + global_offset += m_block_accum[block_idx - 1]; + } + + // Write to sparse output buffers + m_batch_ids[global_offset] = bid; + m_camera_ids[global_offset] = cid; + m_gaussian_ids[global_offset] = gid; + m_radii[global_offset * 2] = (int32_t)radius_x; + m_radii[global_offset * 2 + 1] = (int32_t)radius_y; + m_means2d[global_offset * 2] = mean2d.x; + m_means2d[global_offset * 2 + 1] = mean2d.y; + m_depths[global_offset] = mean_c.z; + m_conics[global_offset * 3] = covar2d_inv[0][0]; + m_conics[global_offset * 3 + 1] = covar2d_inv[0][1]; + m_conics[global_offset * 3 + 2] = covar2d_inv[1][1]; + if (m_compensations != nullptr) { + m_compensations[global_offset] = compensation; + } + } + } + // Lane 0 of the first block in each row writes the indptr. + if (local_id == 0 && block_col_idx == 0) { + if (row_idx == 0) { + m_indptr[0] = 0; + // The final count is written by the host after a scan over + // block_accum. m_indptr[m_B * m_C] = m_block_accum[m_B * + // m_C * blocks_per_row - 1]; + } else { + m_indptr[row_idx] = m_block_accum[block_idx - 1]; + } + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/ProjBwdKernel.hpp b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp new file mode 100644 index 00000000..5198f85f --- /dev/null +++ b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp @@ -0,0 +1,137 @@ +#pragma once + +#include "Common.h" +#include "proj.hpp" + +namespace gsplat::xpu { + +template struct ProjBwdKernel { + + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [C, N, 3] + const T *m_covars; // [C, N, 3, 3] + const T *m_Ks; // [C, 3, 3] + const uint32_t m_width; + const uint32_t m_height; + const CameraModelType m_camera_model; + const T *m_v_means2d; // [C, N, 2] + const T *m_v_covars2d; // [C, N, 2, 2] + T *m_v_means; // [C, N, 3] + T *m_v_covars; // [C, N, 3, 3] + + ProjBwdKernel( + const uint32_t C, + const uint32_t N, + const T *means, + const T *covars, + const T *Ks, + const uint32_t width, + const uint32_t height, + const CameraModelType camera_model, + const T *v_means2d, + const T *v_covars2d, + T *v_means, + T *v_covars + ) + : m_C(C), m_N(N), m_means(means), m_covars(covars), m_Ks(Ks), + m_width(width), m_height(height), m_camera_model(camera_model), + m_v_means2d(v_means2d), m_v_covars2d(v_covars2d), m_v_means(v_means), + m_v_covars(v_covars) {} + + void operator()(sycl::nd_item<1> work_item) const { + + uint32_t idx = work_item.get_global_id(0); + const uint32_t total_gaussians = + (work_item.get_group_range(0) * work_item.get_local_range(0)); + if (idx >= total_gaussians) { + return; + } + + const uint32_t bid = idx / (m_C * m_N); // batch id + const uint32_t cid = (idx / m_N) % m_C; // camera id + + const T *means = m_means + (idx * 3); + const T *covars = m_covars + (idx * 9); + T *v_means = m_v_means + (idx * 3); + T *v_covars = m_v_covars + (idx * 9); + // Correctly index Ks using batch and camera id + const T *Ks = m_Ks + (bid * m_C * 9) + (cid * 9); + const T *v_means2d = m_v_means2d + (idx * 2); + const T *v_covars2d = m_v_covars2d + (idx * 4); + + T fx = Ks[0], cx = Ks[2], fy = Ks[4], cy = Ks[5]; + mat3 v_covar(0.f); + vec3 v_mean(0.f); + const vec3 mean = glm::make_vec3(means); + const mat3 covar = glm::make_mat3(covars); + const vec2 v_mean2d = glm::make_vec2(v_means2d); + const mat2 v_covar2d = glm::make_mat2(v_covars2d); + + switch (m_camera_model) { + case CameraModelType::PINHOLE: // perspective projection + persp_proj_vjp( + mean, + covar, + fx, + fy, + cx, + cy, + m_width, + m_height, + glm::transpose(v_covar2d), + v_mean2d, + v_mean, + v_covar + ); + break; + case CameraModelType::ORTHO: // orthographic projection + ortho_proj_vjp( + mean, + covar, + fx, + fy, + cx, + cy, + m_width, + m_height, + glm::transpose(v_covar2d), + v_mean2d, + v_mean, + v_covar + ); + break; + case CameraModelType::FISHEYE: // fisheye projection + fisheye_proj_vjp( + mean, + covar, + fx, + fy, + cx, + cy, + m_width, + m_height, + glm::transpose(v_covar2d), + v_mean2d, + v_mean, + v_covar + ); + break; + } +// write to outputs: glm is column-major but we want row-major +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + v_covars[i * 3 + j] = T(v_covar[j][i]); + } + } + +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { + v_means[i] = T(v_mean[i]); + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/ProjFwdKernel.hpp b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp new file mode 100644 index 00000000..411e1c01 --- /dev/null +++ b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp @@ -0,0 +1,93 @@ +#pragma once + +#include "Common.h" +#include "proj.hpp" + +namespace gsplat::xpu { + +template struct ProjFwdKernel { + + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [C, N, 3] + const T *m_covars; // [C, N, 3, 3] + const T *m_Ks; // [C, 3, 3] + const uint32_t m_width; + const uint32_t m_height; + const CameraModelType m_camera_model; + T *m_means2d; // [C, N, 2] + T *m_covars2d; // [C, N, 2, 2] + + ProjFwdKernel( + const uint32_t C, + const uint32_t N, + const T *means, // [C, N, 3] + const T *covars, // [C, N, 3, 3] + const T *Ks, // [C, 3, 3] + const uint32_t width, + const uint32_t height, + const CameraModelType camera_model, + T *means2d, // [C, N, 2] + T *covars2d // [C, N, 2, 2] + ) + : m_C(C), m_N(N), m_means(means), m_covars(covars), m_Ks(Ks), + m_width(width), m_height(height), m_camera_model(camera_model), + m_means2d(means2d), m_covars2d(covars2d) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + const uint32_t total_gaussians = + (work_item.get_group_range(0) * work_item.get_local_range(0)); + if (idx >= total_gaussians) { + return; + } + + const uint32_t bid = idx / (m_C * m_N); // batch id + const uint32_t cid = (idx / m_N) % m_C; // camera id + + const T *means = m_means + (idx * 3); + const T *covars = m_covars + (idx * 9); + const T *Ks = m_Ks + (bid * m_C * 9) + (cid * 9); + + T *means2d = m_means2d + (idx * 2); + T *covars2d = m_covars2d + (idx * 4); + + T fx = Ks[0], cx = Ks[2], fy = Ks[4], cy = Ks[5]; + mat2 covar2d(0.f); + vec2 mean2d(0.f); + const vec3 mean = glm::make_vec3(means); + const mat3 covar = glm::make_mat3(covars); + + switch (m_camera_model) { + case CameraModelType::PINHOLE: // perspective projection + persp_proj( + mean, covar, fx, fy, cx, cy, m_width, m_height, covar2d, mean2d + ); + break; + case CameraModelType::ORTHO: // orthographic projection + ortho_proj( + mean, covar, fx, fy, cx, cy, m_width, m_height, covar2d, mean2d + ); + break; + case CameraModelType::FISHEYE: // fisheye projection + fisheye_proj( + mean, covar, fx, fy, cx, cy, m_width, m_height, covar2d, mean2d + ); + break; + } + +#pragma unroll + for (uint32_t i = 0; i < 2; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 2; j++) { // cols + covars2d[i * 2 + j] = T(covar2d[j][i]); + } + } +#pragma unroll + for (uint32_t i = 0; i < 2; i++) { + means2d[i] = T(mean2d[i]); + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp new file mode 100644 index 00000000..d938a035 --- /dev/null +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp @@ -0,0 +1,306 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { + +template inline T sum(vec3 a) { return a.x + a.y + a.z; } + +template +inline void compute_ray_transforms_aabb_vjp( + const T *ray_transforms, + const T *v_means2d, + const vec3 v_normals, + const mat3 W, + const mat3 P, + const vec3 cam_pos, + const vec3 mean_w, + const vec3 mean_c, + const vec4 quat, + const vec2 scale, + mat3 &_v_ray_transforms, + vec4 &v_quat, + vec2 &v_scale, + vec3 &v_mean, + mat3 &v_R, + vec3 &v_t +) { + if (v_means2d[0] != 0 || v_means2d[1] != 0) { + const T distance = ray_transforms[6] * ray_transforms[6] + + ray_transforms[7] * ray_transforms[7] - + ray_transforms[8] * ray_transforms[8]; + const T f = T(1) / (distance); + const T dpx_dT00 = f * ray_transforms[6]; + const T dpx_dT01 = f * ray_transforms[7]; + const T dpx_dT02 = -f * ray_transforms[8]; + const T dpy_dT10 = f * ray_transforms[6]; + const T dpy_dT11 = f * ray_transforms[7]; + const T dpy_dT12 = -f * ray_transforms[8]; + const T dpx_dd = -f * f * + (ray_transforms[0] * ray_transforms[6] + + ray_transforms[1] * ray_transforms[7] - + ray_transforms[2] * ray_transforms[8]); + const T dpx_dT30 = + ray_transforms[0] * f + T(2) * dpx_dd * ray_transforms[6]; + const T dpx_dT31 = + ray_transforms[1] * f + T(2) * dpx_dd * ray_transforms[7]; + const T dpx_dT32 = + -ray_transforms[2] * f - T(2) * dpx_dd * ray_transforms[8]; + const T dpy_dd = -f * f * + (ray_transforms[3] * ray_transforms[6] + + ray_transforms[4] * ray_transforms[7] - + ray_transforms[5] * ray_transforms[8]); + const T dpy_dT30 = + ray_transforms[3] * f + T(2) * dpy_dd * ray_transforms[6]; + const T dpy_dT31 = + ray_transforms[4] * f + T(2) * dpy_dd * ray_transforms[7]; + const T dpy_dT32 = + -ray_transforms[5] * f - T(2) * dpy_dd * ray_transforms[8]; + + _v_ray_transforms[0][0] += v_means2d[0] * dpx_dT00; + _v_ray_transforms[0][1] += v_means2d[0] * dpx_dT01; + _v_ray_transforms[0][2] += v_means2d[0] * dpx_dT02; + _v_ray_transforms[1][0] += v_means2d[1] * dpy_dT10; + _v_ray_transforms[1][1] += v_means2d[1] * dpy_dT11; + _v_ray_transforms[1][2] += v_means2d[1] * dpy_dT12; + _v_ray_transforms[2][0] += + v_means2d[0] * dpx_dT30 + v_means2d[1] * dpy_dT30; + _v_ray_transforms[2][1] += + v_means2d[0] * dpx_dT31 + v_means2d[1] * dpy_dT31; + _v_ray_transforms[2][2] += + v_means2d[0] * dpx_dT32 + v_means2d[1] * dpy_dT32; + } + + mat3 R = quat_to_rotmat(quat); + mat3 v_M = P * glm::transpose(_v_ray_transforms); + mat3 W_t = glm::transpose(W); + mat3 v_RS = W_t * v_M; + vec3 v_tn = W_t * v_normals; + + // dual visible + vec3 tn = W * R[2]; + T cos = glm::dot(-tn, mean_c); + T multiplier = cos > T(0) ? T(1) : T(-1); + v_tn *= multiplier; + + mat3 v_Rot = mat3(v_RS[0] * scale[0], v_RS[1] * scale[1], v_tn); + + quat_to_rotmat_vjp(quat, v_Rot, v_quat); + v_scale[0] += glm::dot(v_RS[0], R[0]); + v_scale[1] += glm::dot(v_RS[1], R[1]); + + v_mean += v_RS[2]; + + v_R += glm::outerProduct(v_M[2], mean_w); + + mat3 RS = quat_to_rotmat(quat) * mat3( + scale[0], + T(0.0), + T(0.0), + T(0.0), + scale[1], + T(0.0), + T(0.0), + T(0.0), + T(1.0) + ); + mat3 v_RS_cam = mat3(v_M[0], v_M[1], v_normals * multiplier); + + v_R += v_RS_cam * glm::transpose(RS); + v_t += v_M[2]; +} + +template struct Projection2DGSFusedBwdKernel { + // fwd inputs + const uint32_t m_B; + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [B, N, 3] + const T *m_quats; // [B, N, 4] + const T *m_scales; // [B, N, 3] + const T *m_viewmats; // [B, C, 4, 4] + const T *m_Ks; // [B, C, 3, 3] + const uint32_t m_image_width; + const uint32_t m_image_height; + // fwd outputs + const int32_t *m_radii; // [B, C, N, 2] + const T *m_ray_transforms; // [B, C, N, 3, 3] + // grad outputs + const T *m_v_means2d; // [B, C, N, 2] + const T *m_v_depths; // [B, C, N] + const T *m_v_normals; // [B, C, N, 3] + const T *m_v_ray_transforms; // [B, C, N, 3, 3] + // grad inputs + T *m_v_means; // [B, N, 3] + T *m_v_quats; // [B, N, 4] + T *m_v_scales; // [B, N, 3] + T *m_v_viewmats; // [B, C, 4, 4] + + Projection2DGSFusedBwdKernel( + const uint32_t B, + const uint32_t C, + const uint32_t N, + const T *means, + const T *quats, + const T *scales, + const T *viewmats, + const T *Ks, + const uint32_t image_width, + const uint32_t image_height, + const int32_t *radii, + const T *ray_transforms, + const T *v_means2d, + const T *v_depths, + const T *v_normals, + const T *v_ray_transforms, + T *v_means, + T *v_quats, + T *v_scales, + T *v_viewmats + ) + : m_B(B), m_C(C), m_N(N), m_means(means), m_quats(quats), + m_scales(scales), m_viewmats(viewmats), m_Ks(Ks), + m_image_width(image_width), m_image_height(image_height), + m_radii(radii), m_ray_transforms(ray_transforms), + m_v_means2d(v_means2d), m_v_depths(v_depths), m_v_normals(v_normals), + m_v_ray_transforms(v_ray_transforms), m_v_means(v_means), + m_v_quats(v_quats), m_v_scales(v_scales), m_v_viewmats(v_viewmats) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + + if (idx >= m_B * m_C * m_N) { + return; + } + + // Check if radii are valid + if (m_radii[idx * 2] <= 0 || m_radii[idx * 2 + 1] <= 0) { + return; + } + + const uint32_t bid = idx / (m_C * m_N); // batch id + const uint32_t cid = (idx / m_N) % m_C; // camera id + const uint32_t gid = idx % m_N; // gaussian id + + // Shift pointers to current camera and gaussian + const T *means = m_means + bid * m_N * 3 + gid * 3; + const T *viewmats = m_viewmats + bid * m_C * 16 + cid * 16; + const T *Ks = m_Ks + bid * m_C * 9 + cid * 9; + + const T *ray_transforms = m_ray_transforms + idx * 9; + + const T *v_means2d = m_v_means2d + idx * 2; + const T *v_depths = m_v_depths + idx; + const T *v_normals = m_v_normals + idx * 3; + const T *v_ray_transforms = m_v_ray_transforms + idx * 9; + + // Transform Gaussian to camera space + mat3 R = mat3( + viewmats[0], + viewmats[4], + viewmats[8], // 1st column + viewmats[1], + viewmats[5], + viewmats[9], // 2nd column + viewmats[2], + viewmats[6], + viewmats[10] // 3rd column + ); + vec3 t = vec3(viewmats[3], viewmats[7], viewmats[11]); + + vec3 mean_w = vec3(means[0], means[1], means[2]); + vec3 mean_c; + pos_world_to_cam(R, t, mean_w, mean_c); + + const T *quats_ptr = m_quats + bid * m_N * 4 + gid * 4; + const T *scales_ptr = m_scales + bid * m_N * 3 + gid * 3; + + vec4 quat = + vec4(quats_ptr[0], quats_ptr[1], quats_ptr[2], quats_ptr[3]); + vec2 scale = vec2(scales_ptr[0], scales_ptr[1]); + + mat3 P = mat3( + Ks[0], T(0.0), Ks[2], T(0.0), Ks[4], Ks[5], T(0.0), T(0.0), T(1.0) + ); + + mat3 _v_ray_transforms = mat3( + v_ray_transforms[0], + v_ray_transforms[1], + v_ray_transforms[2], + v_ray_transforms[3], + v_ray_transforms[4], + v_ray_transforms[5], + v_ray_transforms[6], + v_ray_transforms[7], + v_ray_transforms[8] + ); + + // Add depth gradient to the last element + _v_ray_transforms[2][2] += v_depths[0]; + + vec3 v_normal = vec3(v_normals[0], v_normals[1], v_normals[2]); + + vec3 v_mean = vec3(T(0.0)); + vec2 v_scale = vec2(T(0.0)); + vec4 v_quat = vec4(T(0.0)); + mat3 v_R = mat3(T(0.0)); + vec3 v_t = vec3(T(0.0)); + + // Compute gradients using VJP + compute_ray_transforms_aabb_vjp( + ray_transforms, + v_means2d, + v_normal, + R, + P, + t, + mean_w, + mean_c, + quat, + scale, + _v_ray_transforms, + v_quat, + v_scale, + v_mean, + v_R, + v_t + ); + + // Write out results with atomic additions + if (m_v_means != nullptr) { + T *v_means_out = m_v_means + bid * m_N * 3 + gid * 3; + gpuAtomicAdd(v_means_out, v_mean.x); + gpuAtomicAdd(v_means_out + 1, v_mean.y); + gpuAtomicAdd(v_means_out + 2, v_mean.z); + } + + // Gradients w.r.t. quaternion and scale + T *v_quats_out = m_v_quats + bid * m_N * 4 + gid * 4; + T *v_scales_out = m_v_scales + bid * m_N * 3 + gid * 3; + + gpuAtomicAdd(v_quats_out, v_quat.x); + gpuAtomicAdd(v_quats_out + 1, v_quat.y); + gpuAtomicAdd(v_quats_out + 2, v_quat.z); + gpuAtomicAdd(v_quats_out + 3, v_quat.w); + + gpuAtomicAdd(v_scales_out, v_scale.x); + gpuAtomicAdd(v_scales_out + 1, v_scale.y); + + if (m_v_viewmats != nullptr) { + T *v_viewmats_out = m_v_viewmats + bid * m_C * 16 + cid * 16; + + // Write rotation gradients (column-major to row-major) + for (uint32_t i = 0; i < 3; i++) { + for (uint32_t j = 0; j < 3; j++) { + gpuAtomicAdd(v_viewmats_out + i * 4 + j, v_R[j][i]); + } + gpuAtomicAdd(v_viewmats_out + i * 4 + 3, v_t[i]); + } + } + } +}; + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp new file mode 100644 index 00000000..fb97c166 --- /dev/null +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp @@ -0,0 +1,220 @@ +#pragma once + +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { + +template inline float sum(vec3 a) { return a.x + a.y + a.z; } + +template struct Projection2DGSFusedFwdKernel { + const uint32_t m_B; + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [B, N, 3] + const T *m_quats; // [B, N, 4] + const T *m_scales; // [B, N, 3] + const T *m_viewmats; // [B, C, 4, 4] + const T *m_Ks; // [B, C, 3, 3] + const int32_t m_image_width; + const int32_t m_image_height; + const T m_near_plane; + const T m_far_plane; + const T m_radius_clip; + // outputs + int32_t *m_radii; // [B, C, N, 2] + T *m_means2d; // [B, C, N, 2] + T *m_depths; // [B, C, N] + T *m_ray_transforms; // [B, C, N, 3, 3] + T *m_normals; // [B, C, N, 3] + + Projection2DGSFusedFwdKernel( + const uint32_t B, + const uint32_t C, + const uint32_t N, + const T *means, + const T *quats, + const T *scales, + const T *viewmats, + const T *Ks, + const int32_t image_width, + const int32_t image_height, + const T near_plane, + const T far_plane, + const T radius_clip, + int32_t *radii, + T *means2d, + T *depths, + T *ray_transforms, + T *normals + ) + : m_B(B), m_C(C), m_N(N), m_means(means), m_quats(quats), + m_scales(scales), m_viewmats(viewmats), m_Ks(Ks), + m_image_width(image_width), m_image_height(image_height), + m_near_plane(near_plane), m_far_plane(far_plane), + m_radius_clip(radius_clip), m_radii(radii), m_means2d(means2d), + m_depths(depths), m_ray_transforms(ray_transforms), + m_normals(normals) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + + if (idx >= m_B * m_C * m_N) { + return; + } + + const uint32_t bid = idx / (m_C * m_N); // batch id + const uint32_t cid = (idx / m_N) % m_C; // camera id + const uint32_t gid = idx % m_N; // gaussian id + + // Load data and construct pointers + const T *means = m_means + bid * m_N * 3 + gid * 3; + const T *viewmats = m_viewmats + bid * m_C * 16 + cid * 16; + const T *Ks = m_Ks + bid * m_C * 9 + cid * 9; + + // glm is column-major but input is row-major + // Rotation component of the camera (explicit transpose) + mat3 R = mat3( + viewmats[0], + viewmats[4], + viewmats[8], // 1st column + viewmats[1], + viewmats[5], + viewmats[9], // 2nd column + viewmats[2], + viewmats[6], + viewmats[10] // 3rd column + ); + + // Translation component of the camera + vec3 t = vec3(viewmats[3], viewmats[7], viewmats[11]); + + // Transform Gaussian center to camera space + vec3 mean_c; + pos_world_to_cam(R, t, vec3(means[0], means[1], means[2]), mean_c); + + // Return if primitive is outside valid depth range + if (mean_c.z <= m_near_plane || mean_c.z >= m_far_plane) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + const T *quats = m_quats + bid * m_N * 4 + gid * 4; + const T *scales = m_scales + bid * m_N * 3 + gid * 3; + + // Build rotation matrix from quaternion (quat_to_rotmat returns a mat3) + mat3 rot_mat = + quat_to_rotmat(vec4(quats[0], quats[1], quats[2], quats[3])); + + // Build scale matrix (only x and y for 2D, z is 1) + mat3 scale_mat = mat3( + scales[0], + T(0.0), + T(0.0), + T(0.0), + scales[1], + T(0.0), + T(0.0), + T(0.0), + T(1.0) + ); + + // RS_camera = R * quat_to_rotmat * scale_mat + mat3 RS_camera = R * rot_mat * scale_mat; + + // WH = [RS_camera[0], RS_camera[1], mean_c] + mat3 WH = mat3(RS_camera[0], RS_camera[1], mean_c); + + // Projective transformation matrix: Camera -> Screen + // K^T in column-major order + mat3 world_2_pix = mat3( + Ks[0], T(0.0), Ks[2], T(0.0), Ks[4], Ks[5], T(0.0), T(0.0), T(1.0) + ); + + // M = (WH)^T * K^T + mat3 M = glm::transpose(WH) * world_2_pix; + + // Compute AABB + const vec3 M0 = + vec3(M[0][0], M[0][1], M[0][2]); // first row of KWH + const vec3 M1 = + vec3(M[1][0], M[1][1], M[1][2]); // second row of KWH + const vec3 M2 = + vec3(M[2][0], M[2][1], M[2][2]); // third row of KWH + + const vec3 temp_point = vec3(T(1.0), T(1.0), T(-1.0)); + + // Algebraic manipulation for computing mean and radius + const T distance = sum(temp_point * M2 * M2); + + // Ignore ill-conditioned primitives + if (distance == T(0.0)) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + const vec3 f = (T(1.0) / distance) * temp_point; + const vec2 mean2d = vec2(sum(f * M0 * M2), sum(f * M1 * M2)); + const vec2 temp = {sum(f * M0 * M0), sum(f * M1 * M1)}; + + const vec2 half_extend = mean2d * mean2d - temp; + + const T radius_x = + sycl::ceil(T(3.33) * sycl::sqrt(sycl::max(T(1e-4), half_extend.x))); + const T radius_y = + sycl::ceil(T(3.33) * sycl::sqrt(sycl::max(T(1e-4), half_extend.y))); + + if (radius_x <= m_radius_clip && radius_y <= m_radius_clip) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + // Culling: mask out gaussians outside the image region + if (mean2d.x + radius_x <= T(0) || + mean2d.x - radius_x >= m_image_width || + mean2d.y + radius_y <= T(0) || + mean2d.y - radius_y >= m_image_height) { + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; + return; + } + + // Compute normals (dual visible) + // vec3 normal = vec3(RS_camera[2][0], RS_camera[2][1], + // RS_camera[2][2]); + vec3 normal = RS_camera[2]; + + // Flip normal if it is pointing away from the camera + T multiplier = dot(-normal, mean_c) > T(0) ? T(1.0) : T(-1.0); + normal *= multiplier; + + // Write to outputs + m_radii[idx * 2] = (int32_t)radius_x; + m_radii[idx * 2 + 1] = (int32_t)radius_y; + m_means2d[idx * 2] = mean2d.x; + m_means2d[idx * 2 + 1] = mean2d.y; + m_depths[idx] = mean_c.z; + + // Store ray transforms (row major KWH) + m_ray_transforms[idx * 9 + 0] = M0.x; // [b,c,n,0,0] + m_ray_transforms[idx * 9 + 1] = M0.y; // [b,c,n,0,1] + m_ray_transforms[idx * 9 + 2] = M0.z; // [b,c,n,0,2] + m_ray_transforms[idx * 9 + 3] = M1.x; // [b,c,n,1,0] + m_ray_transforms[idx * 9 + 4] = M1.y; // [b,c,n,1,1] + m_ray_transforms[idx * 9 + 5] = M1.z; // [b,c,n,1,2] + m_ray_transforms[idx * 9 + 6] = M2.x; // [b,c,n,2,0] + m_ray_transforms[idx * 9 + 7] = M2.y; // [b,c,n,2,1] + m_ray_transforms[idx * 9 + 8] = M2.z; // [b,c,n,2,2] + + // Store primitive normals + m_normals[idx * 3] = normal.x; + m_normals[idx * 3 + 1] = normal.y; + m_normals[idx * 3 + 2] = normal.z; + } +}; + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp new file mode 100644 index 00000000..8c7e2b54 --- /dev/null +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp @@ -0,0 +1,114 @@ +#pragma once + +#include "quat_scale_to_covar_preci.hpp" + +namespace gsplat::xpu { + +template struct QuatScaleToCovarPreciBwdKernel { + + const uint32_t m_N; + // fwd inputs + const T *m_quats; // [N, 4] + const T *m_scales; // [N, 3] + // grad outputs + const T *m_v_covars; // [N, 3, 3] or [N, 6] + const T *m_v_precis; // [N, 3, 3] or [N, 6] + const bool m_triu; + // grad inputs + T *m_v_scales; // [N, 3] + T *m_v_quats; // [N, 4] + + QuatScaleToCovarPreciBwdKernel( + const uint32_t N, + const T *quats, + const T *scales, + const T *v_covars, + const T *v_precis, + const bool triu, + T *v_scales, + T *v_quats + ) + : m_N(N), m_quats(quats), m_scales(scales), m_v_covars(v_covars), + m_v_precis(v_precis), m_triu(triu), m_v_scales(v_scales), + m_v_quats(v_quats) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + if (idx >= m_N) { + return; + } + + T *v_scales = m_v_scales + (idx * 3); + T *v_quats = m_v_quats + (idx * 4); + + vec4 quat = glm::make_vec4(m_quats + (idx * 4)); + vec3 scale = glm::make_vec3(m_scales + (idx * 3)); + mat3 rotmat = quat_to_rotmat(quat); + + vec4 v_quat(0.f); + vec3 v_scale(0.f); + + if (m_v_covars != nullptr) { + // glm is column-major, input is row-major + mat3 v_covar; + if (m_triu) { + const T *v_covars = m_v_covars + (idx * 6); + v_covar = mat3( + v_covars[0], + v_covars[1] * .5f, + v_covars[2] * .5f, + v_covars[1] * .5f, + v_covars[3], + v_covars[4] * .5f, + v_covars[2] * .5f, + v_covars[4] * .5f, + v_covars[5] + ); + } else { + const T *v_covars = m_v_covars + (idx * 9); + mat3 v_covar_cast = glm::make_mat3(v_covars); + v_covar = glm::transpose(v_covar_cast); + } + quat_scale_to_covar_vjp( + quat, scale, rotmat, v_covar, v_quat, v_scale + ); + } + + if (m_v_precis != nullptr) { + // glm is column-major, input is row-major + mat3 v_preci; + if (m_triu) { + const T *v_precis = m_v_precis + (idx * 6); + v_preci = mat3( + v_precis[0], + v_precis[1] * .5f, + v_precis[2] * .5f, + v_precis[1] * .5f, + v_precis[3], + v_precis[4] * .5f, + v_precis[2] * .5f, + v_precis[4] * .5f, + v_precis[5] + ); + } else { + const T *v_precis = m_v_precis + (idx * 9); + mat3 v_precis_cast = glm::make_mat3(v_precis); + v_preci = glm::transpose(v_precis_cast); + } + quat_scale_to_preci_vjp( + quat, scale, rotmat, v_preci, v_quat, v_scale + ); + } + +#pragma unroll + for (uint32_t k = 0; k < 3; ++k) { + v_scales[k] = T(v_scale[k]); + } +#pragma unroll + for (uint32_t k = 0; k < 4; ++k) { + v_quats[k] = T(v_quat[k]); + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp new file mode 100644 index 00000000..fef9cc63 --- /dev/null +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp @@ -0,0 +1,92 @@ +#pragma once + +#include "quat_scale_to_covar_preci.hpp" + +namespace gsplat::xpu { + +template struct QuatScaleToCovarPreciFwdKernel { + + const uint32_t m_N; + const T *m_quats; // [N, 4] + const T *m_scales; // [N, 3] + const bool m_triu; + // outputs + T *m_covars; // [N, 3, 3] or [N, 6] + T *m_precis; // [N, 3, 3] or [N, 6] + + QuatScaleToCovarPreciFwdKernel( + const uint32_t N, + const T *quats, + const T *scales, + const bool triu, + T *covars, + T *precis + ) + : m_N(N), m_quats(quats), m_scales(scales), m_triu(triu), + m_covars(covars), m_precis(precis) {} + + void operator()(sycl::nd_item<1> work_item) const { + uint32_t idx = work_item.get_global_id(0); + if (idx >= m_N) { + return; + } + + const T *quats = m_quats + (idx * 4); + const T *scales = m_scales + (idx * 3); + + mat3 covar, preci; + const vec4 quat = glm::make_vec4(quats); + const vec3 scale = glm::make_vec3(scales); + quat_scale_to_covar_preci( + quat, + scale, + m_covars ? &covar : nullptr, + m_precis ? &preci : nullptr + ); + + // write to outputs: glm is column-major but we want row-major + if (m_covars != nullptr) { + if (m_triu) { + T *covars = m_covars + (idx * 6); + covars[0] = T(covar[0][0]); + covars[1] = T(covar[0][1]); + covars[2] = T(covar[0][2]); + covars[3] = T(covar[1][1]); + covars[4] = T(covar[1][2]); + covars[5] = T(covar[2][2]); + } else { + T *covars = m_covars + (idx * 9); +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + covars[i * 3 + j] = T(covar[j][i]); + } + } + } + } + + if (m_precis != nullptr) { + if (m_triu) { + T *precis = m_precis + (idx * 6); + precis[0] = T(preci[0][0]); + precis[1] = T(preci[0][1]); + precis[2] = T(preci[0][2]); + precis[3] = T(preci[1][1]); + precis[4] = T(preci[1][2]); + precis[5] = T(preci[2][2]); + } else { + T *precis = m_precis + (idx * 9); +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + precis[i * 3 + j] = T(preci[j][i]); + } + } + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp new file mode 100644 index 00000000..f00108b6 --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp @@ -0,0 +1,715 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "types.hpp" +#include + +namespace gsplat::xpu { + +// Constants from the CUDA implementation +constexpr float ALPHA_THRESHOLD = 1.0f / 255.0f; +constexpr float FILTER_INV_SQUARE_2DGS = 2.0f; + +template struct RasterizeToPixels2DGSBwdKernel { + // Number of images, gaussians, and intersections + const uint32_t m_I; + const uint32_t m_N; + const uint32_t m_n_isects; + const bool m_packed; + const uint32_t m_chunk_size; + + // Forward pass inputs + const sycl::vec *m_means2d; // Projected Gaussian means + const float *m_ray_transforms; // Transformation matrices + const float *m_colors; // Gaussian colors + const float *m_opacities; // Gaussian opacities + const float *m_normals; // Normals in camera space + const float *m_backgrounds; // Background colors + const bool *m_masks; // Tile masks + + // Image and tile dimensions + const uint32_t m_image_width; + const uint32_t m_image_height; + const uint32_t m_tile_size; + const uint32_t m_tile_width; + const uint32_t m_tile_height; + + // Intersection data + const int32_t *m_tile_offsets; // Intersection offsets + const int32_t *m_flatten_ids; // Global flatten indices + + // Forward pass outputs + const float *m_render_colors; // Rendered colors + const float *m_render_alphas; // Alpha values + const int32_t *m_last_ids; // Last Gaussian indices + const int32_t *m_median_ids; // Median Gaussian indices + + // Gradients from upstream + const float *m_v_render_colors; // Gradients of colors + const float *m_v_render_alphas; // Gradients of alphas + const float *m_v_render_normals; // Gradients of normals + const float *m_v_render_distort; // Gradients of distortion + const float *m_v_render_median; // Gradients of median depth + + // Gradient outputs + sycl::vec + *m_v_means2d_abs; // Gradients of means2d (absolute, can be null) + sycl::vec *m_v_means2d; // Gradients of means2d + float *m_v_ray_transforms; // Gradients of ray transforms + float *m_v_colors; // Gradients of colors + float *m_v_opacities; // Gradients of opacities + float *m_v_normals; // Gradients of normals + float *m_v_densify; // Densification gradients + + // Shared memory + sycl::local_accessor m_slm_id_batch; + sycl::local_accessor, 1> m_slm_xy_opacity; + sycl::local_accessor, 1> m_slm_u_Ms; + sycl::local_accessor, 1> m_slm_v_Ms; + sycl::local_accessor, 1> m_slm_w_Ms; + sycl::local_accessor, 1> m_slm_rgbs; + sycl::local_accessor, 1> m_slm_normals; + + RasterizeToPixels2DGSBwdKernel( + const uint32_t I, + const uint32_t N, + const uint32_t n_isects, + const bool packed, + const uint32_t chunk_size, + // Forward inputs + const sycl::vec *means2d, + const float *ray_transforms, + const float *colors, + const float *opacities, + const float *normals, + const float *backgrounds, + const bool *masks, + // Image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + // Intersections + const int32_t *tile_offsets, + const int32_t *flatten_ids, + // Forward outputs + const float *render_colors, + const float *render_alphas, + const int32_t *last_ids, + const int32_t *median_ids, + // Gradient inputs + const float *v_render_colors, + const float *v_render_alphas, + const float *v_render_normals, + const float *v_render_distort, + const float *v_render_median, + // Gradient outputs + sycl::vec *v_means2d_abs, + sycl::vec *v_means2d, + float *v_ray_transforms, + float *v_colors, + float *v_opacities, + float *v_normals, + float *v_densify, + // Shared memory + sycl::local_accessor slm_id_batch, + sycl::local_accessor, 1> slm_xy_opacity, + sycl::local_accessor, 1> slm_u_Ms, + sycl::local_accessor, 1> slm_v_Ms, + sycl::local_accessor, 1> slm_w_Ms, + sycl::local_accessor, 1> slm_rgbs, + sycl::local_accessor, 1> slm_normals + ) + : m_I(I), m_N(N), m_n_isects(n_isects), m_packed(packed), + m_chunk_size(chunk_size), m_means2d(means2d), + m_ray_transforms(ray_transforms), m_colors(colors), + m_opacities(opacities), m_normals(normals), + m_backgrounds(backgrounds), m_masks(masks), + m_image_width(image_width), m_image_height(image_height), + m_tile_size(tile_size), m_tile_width(tile_width), + m_tile_height(tile_height), m_tile_offsets(tile_offsets), + m_flatten_ids(flatten_ids), m_render_colors(render_colors), + m_render_alphas(render_alphas), m_last_ids(last_ids), + m_median_ids(median_ids), m_v_render_colors(v_render_colors), + m_v_render_alphas(v_render_alphas), + m_v_render_normals(v_render_normals), + m_v_render_distort(v_render_distort), + m_v_render_median(v_render_median), m_v_means2d_abs(v_means2d_abs), + m_v_means2d(v_means2d), m_v_ray_transforms(v_ray_transforms), + m_v_colors(v_colors), m_v_opacities(v_opacities), + m_v_normals(v_normals), m_v_densify(v_densify), + m_slm_id_batch(slm_id_batch), m_slm_xy_opacity(slm_xy_opacity), + m_slm_u_Ms(slm_u_Ms), m_slm_v_Ms(slm_v_Ms), m_slm_w_Ms(slm_w_Ms), + m_slm_rgbs(slm_rgbs), m_slm_normals(slm_normals) {} + + [[intel::reqd_sub_group_size(16)]] + void operator()(sycl::nd_item<3> item) const { + // Map thread and block indices + uint32_t image_id = item.get_group(0); // Block index x -> image_id + uint32_t tile_y = item.get_group(1); // Block index y -> tile_y + uint32_t tile_x = item.get_group(2); // Block index z -> tile_x + uint32_t tile_id = tile_y * m_tile_width + tile_x; + + uint32_t i = tile_y * m_tile_size + item.get_local_id(1); // Pixel y + uint32_t j = tile_x * m_tile_size + item.get_local_id(2); // Pixel x + + // Get pointers to data for current image + const int32_t *tile_offsets_ptr = + m_tile_offsets + image_id * m_tile_height * m_tile_width; + const float *render_alphas_ptr = + m_render_alphas + image_id * m_image_height * m_image_width; + const float *render_colors_ptr = + m_render_colors + + image_id * m_image_height * m_image_width * COLOR_DIM; + + const int32_t *last_ids_ptr = + m_last_ids + image_id * m_image_height * m_image_width; + const int32_t *median_ids_ptr = + m_median_ids + image_id * m_image_height * m_image_width; + + const float *v_render_colors_ptr = + m_v_render_colors + + image_id * m_image_height * m_image_width * COLOR_DIM; + const float *v_render_alphas_ptr = + m_v_render_alphas + image_id * m_image_height * m_image_width; + const float *v_render_normals_ptr = + m_v_render_normals + image_id * m_image_height * m_image_width * 3; + const float *v_render_distort_ptr = nullptr; + if (m_v_render_distort != nullptr) { + v_render_distort_ptr = + m_v_render_distort + image_id * m_image_height * m_image_width; + } + const float *v_render_median_ptr = + m_v_render_median + image_id * m_image_height * m_image_width; + + // Background and mask pointers + const float *backgrounds_ptr = m_backgrounds; + if (backgrounds_ptr != nullptr) { + backgrounds_ptr += image_id * COLOR_DIM; + } + + const bool *masks_ptr = m_masks; + if (masks_ptr != nullptr) { + masks_ptr += image_id * m_tile_height * m_tile_width; + } + + // If tile is masked, do nothing + if (masks_ptr != nullptr && !masks_ptr[tile_id]) { + return; + } + + // Pixel center coordinates + const float px = static_cast(j) + 0.5f; + const float py = static_cast(i) + 0.5f; + const int32_t pix_id = static_cast(sycl::min( + static_cast(i * m_image_width + j), + static_cast(m_image_width * m_image_height - 1) + )); + + // Check if pixel is inside image bounds + bool inside = (i < m_image_height && j < m_image_width); + + // Find range of gaussians for this tile + int32_t range_start = tile_offsets_ptr[tile_id]; + int32_t range_end; + if ((image_id == m_I - 1) && + (tile_id == static_cast(m_tile_width * m_tile_height - 1) + )) { + range_end = m_n_isects; + } else { + range_end = tile_offsets_ptr[tile_id + 1]; + } + + // Calculate number of batches needed + uint32_t num_batches = + (range_end - range_start + m_chunk_size - 1) / m_chunk_size; + + // Transmittance after last gaussian + float T_final = 1.0f - render_alphas_ptr[pix_id]; + float T = T_final; + + // Buffers for accumulating contributions + float buffer[COLOR_DIM] = {0.0f}; + float buffer_normals[3] = {0.0f}; + + // Index of last gaussian that contributed to this pixel + const int32_t bin_final = inside ? last_ids_ptr[pix_id] : 0; + + // Index of gaussian that contributes to median depth + const int32_t median_idx = inside ? median_ids_ptr[pix_id] : 0; + + // Get thread rank for shared memory access + uint32_t tr = item.get_local_linear_id(); + + // Load gradients for this pixel + BufferType_t v_render_c{}; + if (inside) { + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + v_render_c = + *reinterpret_cast *>( + v_render_colors_ptr + pix_id * COLOR_DIM + ); + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + v_render_c[k] = v_render_colors_ptr[pix_id * COLOR_DIM + k]; + } + } + } + + float v_render_a = inside ? v_render_alphas_ptr[pix_id] : 0.0f; + + sycl::vec v_render_n{0.0f, 0.0f, 0.0f}; + if (inside) { + v_render_n.x() = v_render_normals_ptr[pix_id * 3]; + v_render_n.y() = v_render_normals_ptr[pix_id * 3 + 1]; + v_render_n.z() = v_render_normals_ptr[pix_id * 3 + 2]; + } + + // Prepare for distortion (if needed) + float v_distort = 0.0f; + float accum_d = 0.0f, accum_w = 0.0f; + float accum_d_buffer = 0.0f, accum_w_buffer = 0.0f, + distort_buffer = 0.0f; + if (v_render_distort_ptr != nullptr && inside) { + v_distort = v_render_distort_ptr[pix_id]; + accum_d_buffer = + render_colors_ptr[pix_id * COLOR_DIM + COLOR_DIM - 1]; + accum_d = accum_d_buffer; + accum_w_buffer = render_alphas_ptr[pix_id]; + accum_w = accum_w_buffer; + } + + // Get median depth gradient + float v_median = inside ? v_render_median_ptr[pix_id] : 0.0f; + + // Find the maximum final gaussian id in the warp + int32_t warp_bin_final = sycl::reduce_over_group( + item.get_sub_group(), bin_final, sycl::maximum() + ); + + // Process batches of gaussians in reverse order (back to front) + for (int32_t b = 0; b < num_batches; ++b) { + // Synchronize threads before loading next batch + item.barrier(sycl::access::fence_space::local_space); + + // Compute batch boundaries + int32_t batch_end = range_end - 1 - m_chunk_size * b; + int32_t batch_size = + sycl::min(m_chunk_size, batch_end + 1 - range_start); + + // Load gaussian data into shared memory (in reverse order) + int32_t idx = batch_end - tr; + + if (idx >= range_start && tr < m_chunk_size) { + int32_t g = m_flatten_ids[idx]; + m_slm_id_batch[tr] = g; + + // Load position and opacity + sycl::vec xy = m_means2d[g]; + float opac = m_opacities[g]; + m_slm_xy_opacity[tr] = sycl::vec(xy[0], xy[1], opac); + + // Load ray transform matrix rows + m_slm_u_Ms[tr] = sycl::vec( + m_ray_transforms[g * 9], + m_ray_transforms[g * 9 + 1], + m_ray_transforms[g * 9 + 2] + ); + m_slm_v_Ms[tr] = sycl::vec( + m_ray_transforms[g * 9 + 3], + m_ray_transforms[g * 9 + 4], + m_ray_transforms[g * 9 + 5] + ); + m_slm_w_Ms[tr] = sycl::vec( + m_ray_transforms[g * 9 + 6], + m_ray_transforms[g * 9 + 7], + m_ray_transforms[g * 9 + 8] + ); + + // Load colors + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + m_slm_rgbs[tr] = *reinterpret_cast< + const BufferType_t *>( + m_colors + g * COLOR_DIM + ); + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + m_slm_rgbs[tr][k] = m_colors[g * COLOR_DIM + k]; + } + } + + // Load normals + m_slm_normals[tr] = sycl::vec( + m_normals[g * 3], m_normals[g * 3 + 1], m_normals[g * 3 + 2] + ); + } + + // Wait for all threads to load data + item.barrier(sycl::access::fence_space::local_space); + + // Process gaussians in batch from back to front + for (int32_t t = sycl::max(0, batch_end - warp_bin_final); + t < batch_size; + ++t) { + bool valid = inside; + if (batch_end - t > bin_final) { + valid = false; + } + + // Variables for forward pass calculations + float alpha = 0.0f, opac = 0.0f, vis = 0.0f; + float gauss_weight_3d = 0.0f, gauss_weight_2d = 0.0f, + gauss_weight = 0.0f; + sycl::vec s{0.0f, 0.0f}, d{0.0f, 0.0f}; + sycl::vec h_u{0.0f, 0.0f, 0.0f}, + h_v{0.0f, 0.0f, 0.0f}; + sycl::vec ray_cross{0.0f, 0.0f, 0.0f}, + w_M{0.0f, 0.0f, 0.0f}; + + // Perform forward pass calculations for current gaussian + if (valid) { + // Get gaussian parameters from shared memory + sycl::vec xy_opac = m_slm_xy_opacity[t]; + opac = xy_opac[2]; + + sycl::vec u_M = m_slm_u_Ms[t]; + sycl::vec v_M = m_slm_v_Ms[t]; + w_M = m_slm_w_Ms[t]; + + // Calculate homogeneous plane parameters + h_u = sycl::vec( + px * w_M[0] - u_M[0], + px * w_M[1] - u_M[1], + px * w_M[2] - u_M[2] + ); + + h_v = sycl::vec( + py * w_M[0] - v_M[0], + py * w_M[1] - v_M[1], + py * w_M[2] - v_M[2] + ); + + // Compute ray intersection using cross product + ray_cross = sycl::cross(h_u, h_v); + + // Check for valid intersection + if (ray_cross[2] == 0.0f) { + valid = false; + } else { + // Project to UV space + s = sycl::vec( + ray_cross[0] / ray_cross[2], + ray_cross[1] / ray_cross[2] + ); + + // Calculate 3D gaussian weight + gauss_weight_3d = s[0] * s[0] + s[1] * s[1]; + + // Calculate 2D projected gaussian weight + d = sycl::vec( + xy_opac[0] - px, xy_opac[1] - py + ); + gauss_weight_2d = FILTER_INV_SQUARE_2DGS * + (d[0] * d[0] + d[1] * d[1]); + + // Use minimum of 3D and 2D weights + gauss_weight = + sycl::min(gauss_weight_3d, gauss_weight_2d); + + // Calculate sigma and alpha + float sigma = 0.5f * gauss_weight; + vis = sycl::exp(-sigma); + alpha = sycl::min(0.999f, opac * vis); + + // Skip if gaussian is transparent + if (sigma < 0.0f || alpha < ALPHA_THRESHOLD) { + valid = false; + } + } + } + + // Skip if no thread in the sub-group has a valid gaussian + bool any_valid = + sycl::any_of_group(item.get_sub_group(), valid); + if (!any_valid) { + continue; + } + + // Initialize gradient variables + BufferType_t v_rgb_local{}; + sycl::vec v_normal_local{0.0f, 0.0f, 0.0f}; + sycl::vec v_u_M_local{0.0f, 0.0f, 0.0f}; + sycl::vec v_v_M_local{0.0f, 0.0f, 0.0f}; + sycl::vec v_w_M_local{0.0f, 0.0f, 0.0f}; + sycl::vec v_xy_local{0.0f, 0.0f}; + sycl::vec v_xy_abs_local{0.0f, 0.0f}; + float v_opacity_local = 0.0f; + + if (valid) { + // Gradient contribution from median depth + if (batch_end - t == median_idx) { + v_rgb_local[COLOR_DIM - 1] += v_median; + } + + // Compute the current T for this gaussian + float ra = 1.0f / (1.0f - alpha); + T *= ra; + + // Weight for the current gaussian + float fac = alpha * T; + + // Update rgb gradients + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + v_rgb_local[k] += fac * v_render_c[k]; + } + + // Calculate alpha gradient + float v_alpha = 0.0f; + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + v_alpha += (m_slm_rgbs[t][k] * T - buffer[k] * ra) * + v_render_c[k]; + } + + // Update normal gradients + for (uint32_t k = 0; k < 3; ++k) { + v_normal_local[k] = fac * v_render_n[k]; + } + + for (uint32_t k = 0; k < 3; ++k) { + v_alpha += + (m_slm_normals[t][k] * T - buffer_normals[k] * ra) * + v_render_n[k]; + } + + // Gradient contribution from alpha + v_alpha += T_final * ra * v_render_a; + + // Adjust alpha gradients by background color + if (backgrounds_ptr != nullptr) { + float accum = 0.0f; + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + accum += backgrounds_ptr[k] * v_render_c[k]; + } + v_alpha += -T_final * ra * accum; + } + + // Contribution from distortion + if (v_render_distort_ptr != nullptr) { + float depth = m_slm_rgbs[t][COLOR_DIM - 1]; + float dl_dw = + 2.0f * + (2.0f * (depth * accum_w_buffer - accum_d_buffer) + + (accum_d - depth * accum_w)); + v_alpha += + (dl_dw * T - distort_buffer * ra) * v_distort; + accum_d_buffer -= fac * depth; + accum_w_buffer -= fac; + distort_buffer += dl_dw * fac; + v_rgb_local[COLOR_DIM - 1] += + 2.0f * fac * (2.0f - 2.0f * T - accum_w + fac) * + v_distort; + } + + // Calculate geometry-related gradients + if (opac * vis <= 0.999f) { + float v_depth = 0.0f; + float v_G = opac * v_alpha; + + // Case 1: Ray-primitive intersection used in forward + // pass + if (gauss_weight_3d <= gauss_weight_2d) { + sycl::vec v_s( + v_G * -vis * s[0] + v_depth * w_M[0], + v_G * -vis * s[1] + v_depth * w_M[1] + ); + + // Backward through projective transform + sycl::vec v_z_w_M(s[0], s[1], 1.0f); + float v_sx_pz = v_s[0] / ray_cross[2]; + float v_sy_pz = v_s[1] / ray_cross[2]; + sycl::vec v_ray_cross( + v_sx_pz, + v_sy_pz, + -(v_sx_pz * s[0] + v_sy_pz * s[1]) + ); + + // Calculate cross products for gradient computation + sycl::vec v_h_u = + sycl::cross(h_v, v_ray_cross); + sycl::vec v_h_v = + sycl::cross(v_ray_cross, h_u); + + // Compute gradients for transformation matrices + v_u_M_local = sycl::vec( + -v_h_u[0], -v_h_u[1], -v_h_u[2] + ); + v_v_M_local = sycl::vec( + -v_h_v[0], -v_h_v[1], -v_h_v[2] + ); + v_w_M_local = sycl::vec( + px * v_h_u[0] + py * v_h_v[0] + + v_depth * v_z_w_M[0], + px * v_h_u[1] + py * v_h_v[1] + + v_depth * v_z_w_M[1], + px * v_h_u[2] + py * v_h_v[2] + + v_depth * v_z_w_M[2] + ); + + // Case 2: 2D projected gaussian used in forward + // pass + } else { + float v_G_ddelx = + -vis * FILTER_INV_SQUARE_2DGS * d[0]; + float v_G_ddely = + -vis * FILTER_INV_SQUARE_2DGS * d[1]; + v_xy_local = sycl::vec( + v_G * v_G_ddelx, v_G * v_G_ddely + ); + + if (m_v_means2d_abs != nullptr) { + v_xy_abs_local = sycl::vec( + sycl::fabs(v_xy_local[0]), + sycl::fabs(v_xy_local[1]) + ); + } + } + + v_opacity_local = vis * v_alpha; + } + + // Update cumulative buffers + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + buffer[k] += m_slm_rgbs[t][k] * fac; + } + + for (uint32_t k = 0; k < 3; ++k) { + buffer_normals[k] += m_slm_normals[t][k] * fac; + } + } + + // Sub-group reduction to sum gradients + auto sub_group = item.get_sub_group(); + + // Reduce RGB gradients + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + v_rgb_local = sycl::reduce_over_group( + sub_group, + v_rgb_local, + sycl::plus>() + ); + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + v_rgb_local[k] = sycl::reduce_over_group( + sub_group, v_rgb_local[k], sycl::plus() + ); + } + } + + // Reduce other gradients + v_normal_local = sycl::reduce_over_group( + sub_group, v_normal_local, sycl::plus>() + ); + v_u_M_local = sycl::reduce_over_group( + sub_group, v_u_M_local, sycl::plus>() + ); + v_v_M_local = sycl::reduce_over_group( + sub_group, v_v_M_local, sycl::plus>() + ); + v_w_M_local = sycl::reduce_over_group( + sub_group, v_w_M_local, sycl::plus>() + ); + v_xy_local = sycl::reduce_over_group( + sub_group, v_xy_local, sycl::plus>() + ); + v_opacity_local = sycl::reduce_over_group( + sub_group, v_opacity_local, sycl::plus() + ); + + if (m_v_means2d_abs != nullptr) { + v_xy_abs_local = sycl::reduce_over_group( + sub_group, + v_xy_abs_local, + sycl::plus>() + ); + } + + // Write gradients to global memory + int32_t g = m_slm_id_batch[t]; + + if (sub_group.get_local_id() == 0) { + // Update color gradients + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + gpuAtomicAddGlobal( + m_v_colors[g * COLOR_DIM + k], v_rgb_local[k] + ); + } + + // Update normal gradients + for (uint32_t k = 0; k < 3; ++k) { + gpuAtomicAddGlobal( + m_v_normals[g * 3 + k], v_normal_local[k] + ); + } + + // Update ray transform gradients + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9], v_u_M_local[0] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 1], v_u_M_local[1] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 2], v_u_M_local[2] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 3], v_v_M_local[0] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 4], v_v_M_local[1] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 5], v_v_M_local[2] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 6], v_w_M_local[0] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 7], v_w_M_local[1] + ); + gpuAtomicAddGlobal( + m_v_ray_transforms[g * 9 + 8], v_w_M_local[2] + ); + + // Update means2d gradients + gpuAtomicAddGlobal(m_v_means2d[g].x(), v_xy_local[0]); + gpuAtomicAddGlobal(m_v_means2d[g].y(), v_xy_local[1]); + + if (m_v_means2d_abs != nullptr) { + gpuAtomicAddGlobal( + m_v_means2d_abs[g].x(), v_xy_abs_local[0] + ); + gpuAtomicAddGlobal( + m_v_means2d_abs[g].y(), v_xy_abs_local[1] + ); + } + + // Update opacity gradients + gpuAtomicAddGlobal(m_v_opacities[g], v_opacity_local); + } + + if (valid) { + float depth = m_slm_w_Ms[t][2]; + m_v_densify[g * 2] = m_v_ray_transforms[g * 9 + 2] * depth; + m_v_densify[g * 2 + 1] = + m_v_ray_transforms[g * 9 + 5] * depth; + } + } + } + } +}; + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp new file mode 100644 index 00000000..bbf5201c --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp @@ -0,0 +1,412 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "types.hpp" +#include + +namespace gsplat::xpu { + +// Constants from the CUDA implementation +constexpr float ALPHA_THRESHOLD = 1.0f / 255.0f; +constexpr float FILTER_INV_SQUARE_2DGS = 2.0f; + +template struct RasterizeToPixels2DGSFwdKernel { + const uint32_t m_I; // number of images + const uint32_t m_N; // number of gaussians + const uint32_t m_n_isects; // number of intersections + const bool m_packed; // whether tensors are packed + const uint32_t m_chunk_size; // chunk size for batch processing + + const sycl::vec *m_means2d; // Projected Gaussian means + const float *m_ray_transforms; // Transformation matrices + const float *m_colors; // Gaussian colors + const float *m_opacities; // Gaussian opacities + const float *m_normals; // Normals in camera space + const float *m_backgrounds; // Background colors + const bool *m_masks; // Tile masks + + const uint32_t m_image_width; + const uint32_t m_image_height; + const uint32_t m_tile_size; + const uint32_t m_tile_width; + const uint32_t m_tile_height; + + const int32_t *m_tile_offsets; // Intersection offsets + const int32_t *m_flatten_ids; // Global flatten indices + + float *m_render_colors; // Output rendered colors + float *m_render_alphas; // Output alpha values + float *m_render_normals; // Output rendered normals + float *m_render_distort; // Output distortion values + float *m_render_median; // Output median depth values + int32_t *m_last_ids; // Output indices of last Gaussians + int32_t *m_median_ids; // Output indices of median Gaussians + + // Shared memory accessors + sycl::local_accessor m_slm_id_batch; + sycl::local_accessor, 1> m_slm_xy_opacity; + sycl::local_accessor, 1> m_slm_u_Ms; + sycl::local_accessor, 1> m_slm_v_Ms; + sycl::local_accessor, 1> m_slm_w_Ms; + + RasterizeToPixels2DGSFwdKernel( + const uint32_t I, + const uint32_t N, + const uint32_t n_isects, + const bool packed, + const uint32_t chunk_size, + const sycl::vec *means2d, + const float *ray_transforms, + const float *colors, + const float *opacities, + const float *normals, + const float *backgrounds, + const bool *masks, + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + const int32_t *tile_offsets, + const int32_t *flatten_ids, + float *render_colors, + float *render_alphas, + float *render_normals, + float *render_distort, + float *render_median, + int32_t *last_ids, + int32_t *median_ids, + sycl::local_accessor slm_id_batch, + sycl::local_accessor, 1> slm_xy_opacity, + sycl::local_accessor, 1> slm_u_Ms, + sycl::local_accessor, 1> slm_v_Ms, + sycl::local_accessor, 1> slm_w_Ms + ) + : m_I(I), m_N(N), m_n_isects(n_isects), m_packed(packed), + m_chunk_size(chunk_size), m_means2d(means2d), + m_ray_transforms(ray_transforms), m_colors(colors), + m_opacities(opacities), m_normals(normals), + m_backgrounds(backgrounds), m_masks(masks), + m_image_width(image_width), m_image_height(image_height), + m_tile_size(tile_size), m_tile_width(tile_width), + m_tile_height(tile_height), m_tile_offsets(tile_offsets), + m_flatten_ids(flatten_ids), m_render_colors(render_colors), + m_render_alphas(render_alphas), m_render_normals(render_normals), + m_render_distort(render_distort), m_render_median(render_median), + m_last_ids(last_ids), m_median_ids(median_ids), + m_slm_id_batch(slm_id_batch), m_slm_xy_opacity(slm_xy_opacity), + m_slm_u_Ms(slm_u_Ms), m_slm_v_Ms(slm_v_Ms), m_slm_w_Ms(slm_w_Ms) {} + + [[intel::reqd_sub_group_size(16)]] + void operator()(sycl::nd_item<3> item) const { + // Map thread and block indices to image, tile, and pixel coordinates + int32_t image_id = item.get_group(0); // Block index x -> image_id + int32_t tile_y = item.get_group(1); // Block index y -> tile_y + int32_t tile_x = item.get_group(2); // Block index z -> tile_x + int32_t tile_id = tile_y * m_tile_width + tile_x; + + uint32_t i = tile_y * m_tile_size + item.get_local_id(1); // Pixel y + uint32_t j = tile_x * m_tile_size + item.get_local_id(2); // Pixel x + + // Get pointers to data for current image + const int32_t *tile_offsets_ptr = + m_tile_offsets + image_id * m_tile_height * m_tile_width; + float *render_colors_ptr = m_render_colors + image_id * m_image_height * + m_image_width * + COLOR_DIM; + float *render_alphas_ptr = + m_render_alphas + image_id * m_image_height * m_image_width; + int32_t *last_ids_ptr = + m_last_ids + image_id * m_image_height * m_image_width; + float *render_normals_ptr = + m_render_normals + image_id * m_image_height * m_image_width * 3; + float *render_distort_ptr = + m_render_distort + image_id * m_image_height * m_image_width; + float *render_median_ptr = + m_render_median + image_id * m_image_height * m_image_width; + int32_t *median_ids_ptr = + m_median_ids + image_id * m_image_height * m_image_width; + + // Background and mask pointers + const float *backgrounds_ptr = m_backgrounds; + if (backgrounds_ptr != nullptr) { + backgrounds_ptr += image_id * COLOR_DIM; + } + + const bool *masks_ptr = m_masks; + if (masks_ptr != nullptr) { + masks_ptr += image_id * m_tile_height * m_tile_width; + } + + // Find pixel center + float px = static_cast(j) + 0.5f; + float py = static_cast(i) + 0.5f; + int32_t pix_id = i * m_image_width + j; + + // Check if pixel is inside image bounds + bool inside = (i < m_image_height && j < m_image_width); + bool done = !inside; + + // Handle masked tiles + if (masks_ptr != nullptr && inside && !masks_ptr[tile_id]) { + // Render background for masked tiles + if (inside) { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + render_colors_ptr[pix_id * COLOR_DIM + k] = + backgrounds_ptr == nullptr ? 0.0f : backgrounds_ptr[k]; + } + } + return; + } + + // Get range of gaussians for this tile + int32_t range_start = tile_offsets_ptr[tile_id]; + int32_t range_end = + (image_id == m_I - 1) && + (tile_id == + static_cast(m_tile_width * m_tile_height - 1)) + ? m_n_isects + : tile_offsets_ptr[tile_id + 1]; + + // Calculate number of batches needed + uint32_t num_batches = + (range_end - range_start + m_chunk_size - 1) / m_chunk_size; + + // Initialize rendering accumulators + float T = 1.0f; // Transmittance + BufferType_t pix_out{}; // Accumulated color + float normal_out[3] = {0.0f}; // Accumulated normal + uint32_t cur_idx = 0; // Current index + float distort = 0.0f; // Distortion + float accum_vis_depth = 0.0f; // Accumulated visibility * depth + float median_depth = 0.0f; // Median depth + uint32_t median_idx = 0; // Median index + + // Get thread rank for shared memory access + uint32_t tr = item.get_local_id(1) * m_tile_size + item.get_local_id(2); + + // Process batches of gaussians + for (uint32_t b = 0; b < num_batches; ++b) { + // Synchronize threads + item.barrier(sycl::access::fence_space::local_space); + + // Each thread loads one gaussian + uint32_t batch_start = range_start + m_chunk_size * b; + uint32_t idx = batch_start + tr; + + if (tr < m_chunk_size && idx < range_end) { + // Get gaussian index + int32_t g = m_flatten_ids[idx]; + m_slm_id_batch[tr] = g; + + // Load gaussian parameters + sycl::vec xy = m_means2d[g]; + float opac = m_opacities[g]; + m_slm_xy_opacity[tr] = sycl::vec(xy[0], xy[1], opac); + + // Load ray transformation matrix rows + m_slm_u_Ms[tr] = sycl::vec( + m_ray_transforms[g * 9], + m_ray_transforms[g * 9 + 1], + m_ray_transforms[g * 9 + 2] + ); + m_slm_v_Ms[tr] = sycl::vec( + m_ray_transforms[g * 9 + 3], + m_ray_transforms[g * 9 + 4], + m_ray_transforms[g * 9 + 5] + ); + m_slm_w_Ms[tr] = sycl::vec( + m_ray_transforms[g * 9 + 6], + m_ray_transforms[g * 9 + 7], + m_ray_transforms[g * 9 + 8] + ); + } + + // Wait for all threads to load data + item.barrier(sycl::access::fence_space::local_space); + + // Manual check for all threads done (instead of CUDA's + // __syncthreads_count) In SYCL, we have to use barrier + // synchronization and local variables for this + + // Process gaussians in the current batch + uint32_t batch_size = + sycl::min(m_chunk_size, range_end - batch_start); + for (uint32_t t = 0; t < batch_size && !done; ++t) { + // Get gaussian parameters from shared memory + const sycl::vec xy_opac = m_slm_xy_opacity[t]; + const float opac = xy_opac[2]; + + // Get transformation matrix rows + const sycl::vec u_M = m_slm_u_Ms[t]; + const sycl::vec v_M = m_slm_v_Ms[t]; + const sycl::vec w_M = m_slm_w_Ms[t]; + + // Calculate homogeneous plane parameters + // h_u = px * w_M - u_M + sycl::vec h_u( + px * w_M[0] - u_M[0], + px * w_M[1] - u_M[1], + px * w_M[2] - u_M[2] + ); + + // h_v = py * w_M - v_M + sycl::vec h_v( + py * w_M[0] - v_M[0], + py * w_M[1] - v_M[1], + py * w_M[2] - v_M[2] + ); + + // Compute intersection using cross product + // ray_cross = h_u × h_v + sycl::vec ray_cross( + h_u[1] * h_v[2] - h_u[2] * h_v[1], + h_u[2] * h_v[0] - h_u[0] * h_v[2], + h_u[0] * h_v[1] - h_u[1] * h_v[0] + ); + + if (ray_cross[2] == 0.0f) { + continue; + } + + // Project to UV space + // s = [ray_cross.x / ray_cross.z, ray_cross.y / ray_cross.z] + sycl::vec s( + ray_cross[0] / ray_cross[2], ray_cross[1] / ray_cross[2] + ); + + // Calculate gaussian weight in 3D + // gauss_weight_3d = s.x * s.x + s.y * s.y + float gauss_weight_3d = s[0] * s[0] + s[1] * s[1]; + + // Calculate projected gaussian weight in 2D + // d = [xy_opac.x - px, xy_opac.y - py] + sycl::vec d(xy_opac[0] - px, xy_opac[1] - py); + // gauss_weight_2d = FILTER_INV_SQUARE_2DGS * (d.x * d.x + d.y * + // d.y) + float gauss_weight_2d = + FILTER_INV_SQUARE_2DGS * (d[0] * d[0] + d[1] * d[1]); + + // Use minimum of 3D and 2D gaussian weights + // gauss_weight = min(gauss_weight_3d, gauss_weight_2d) + float gauss_weight = + sycl::min(gauss_weight_3d, gauss_weight_2d); + + // Calculate sigma and alpha + float sigma = 0.5f * gauss_weight; + float alpha = sycl::min(0.999f, opac * sycl::exp(-sigma)); + + // Skip transparent gaussians + if (sigma < 0.0f || alpha < ALPHA_THRESHOLD) { + continue; + } + + // Calculate next transmittance + float next_T = T * (1.0f - alpha); + if (next_T <= 1e-4f) { + done = true; + break; + } + + // Perform volumetric rendering + int32_t g = m_slm_id_batch[t]; + float vis = alpha * T; + + // Accumulate color + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + const auto *c_ptr = reinterpret_cast< + const BufferType_t *>( + m_colors + g * COLOR_DIM + ); + pix_out += (*c_ptr) * vis; + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + pix_out[k] += m_colors[g * COLOR_DIM + k] * vis; + } + } + + // Accumulate normal + const float *n_ptr = m_normals + g * 3; + for (uint32_t k = 0; k < 3; ++k) { + normal_out[k] += n_ptr[k] * vis; + } + + // Calculate distortion if needed + if (m_render_distort != nullptr) { + const float depth = m_colors[g * COLOR_DIM + COLOR_DIM - 1]; + const float distort_bi_0 = vis * depth * (1.0f - T); + const float distort_bi_1 = vis * accum_vis_depth; + distort += 2.0f * (distort_bi_0 - distort_bi_1); + accum_vis_depth += vis * depth; + } + + // Track median depth + if (T > 0.5f) { + median_depth = m_colors[g * COLOR_DIM + COLOR_DIM - 1]; + median_idx = batch_start + t; + } + + cur_idx = batch_start + t; + T = next_T; + } + } + + // Write results if pixel is inside the image + if (inside) { + // Store alpha (1 - transmittance) + render_alphas_ptr[pix_id] = 1.0f - T; + + // Store color (accumulated + background * transmittance) + if (backgrounds_ptr == nullptr) { + // No background + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + *reinterpret_cast *>( + render_colors_ptr + pix_id * COLOR_DIM + ) = pix_out; + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + render_colors_ptr[pix_id * COLOR_DIM + k] = pix_out[k]; + } + } + } else { + // With background + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + BufferType_t bg; + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + bg[k] = backgrounds_ptr[k]; + } + *reinterpret_cast *>( + render_colors_ptr + pix_id * COLOR_DIM + ) = pix_out + bg * T; + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + render_colors_ptr[pix_id * COLOR_DIM + k] = + pix_out[k] + T * backgrounds_ptr[k]; + } + } + } + + // Store normal + for (uint32_t k = 0; k < 3; ++k) { + render_normals_ptr[pix_id * 3 + k] = normal_out[k]; + } + + // Store last gaussian index + last_ids_ptr[pix_id] = static_cast(cur_idx); + + // Store distortion if needed + if (m_render_distort != nullptr) { + render_distort_ptr[pix_id] = distort; + } + + // Store median depth and index + render_median_ptr[pix_id] = median_depth; + median_ids_ptr[pix_id] = static_cast(median_idx); + } + } +}; + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp new file mode 100644 index 00000000..94fc968e --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp @@ -0,0 +1,447 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "types.hpp" +#include + +namespace gsplat::xpu { + +template +struct RasterizeToPixelsBwdKernel { + // Inputs (fwd inputs) + const uint32_t m_C; + const uint32_t m_N; + const uint32_t m_n_isects; + const bool m_packed; + const uint32_t m_concat_stride; + const S *m_concatenated_data; + const sycl::vec *m_means2d; // [C, N, 2] or [nnz, 2] + const vec3 *m_conics; // [C, N, 3] or [nnz, 3] + const S *m_colors; // [C, N, COLOR_DIM] or [nnz, COLOR_DIM] + const S *m_opacities; // [C, N] or [nnz] + const S *m_backgrounds; // [C, COLOR_DIM] or [nnz, COLOR_DIM] + const bool *m_masks; // [C, tile_height, tile_width] + const uint32_t m_image_width; + const uint32_t m_image_height; + const uint32_t m_tile_size; + const uint32_t m_tile_width; + const uint32_t m_tile_height; + const int32_t *m_tile_offsets; // [C, tile_height, tile_width] + const int32_t *m_flatten_ids; // [n_isects] + + // Forward outputs + const S *m_render_alphas; // [C, image_height, image_width] + const int32_t *m_last_ids; // [C, image_height, image_width] + + // Gradients from downstream (grad outputs) + const S *m_v_render_colors; // [C, image_height, image_width, COLOR_DIM] + const S *m_v_render_alphas; // [C, image_height, image_width] + + // Gradients to be accumulated (grad inputs) + sycl::vec *m_v_means2d_abs; // [C, N, 2] or [nnz, 2] (can be nullptr) + sycl::vec *m_v_means2d; // [C, N, 2] or [nnz, 2] + vec3 *m_v_conics; // [C, N, 3] or [nnz, 3] + S *m_v_colors; // [C, N, COLOR_DIM] or [nnz, COLOR_DIM] + S *m_v_opacities; // [C, N] or [nnz] + + sycl::local_accessor m_slm_flatten_ids; + sycl::local_accessor, 1> m_slm_means2d; + sycl::local_accessor m_slm_opacities; + sycl::local_accessor, 1> m_slm_conics; + sycl::local_accessor, 1> m_slm_colors; + + RasterizeToPixelsBwdKernel( + const uint32_t C, + const uint32_t N, + const uint32_t n_isects, + const bool packed, + const uint32_t concat_stride, + const S *concatenated_data, + const sycl::vec *means2d, + const vec3 *conics, + const S *colors, + const S *opacities, + const S *backgrounds, + const bool *masks, + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + const int32_t *tile_offsets, + const int32_t *flatten_ids, + const S *render_alphas, + const int32_t *last_ids, + const S *v_render_colors, + const S *v_render_alphas, + sycl::vec *v_means2d_abs, + sycl::vec *v_means2d, + vec3 *v_conics, + S *v_colors, + S *v_opacities, + sycl::local_accessor slm_flatten_ids, + sycl::local_accessor, 1> slm_means2d, + sycl::local_accessor slm_opacities, + sycl::local_accessor, 1> slm_conics, + sycl::local_accessor, 1> slm_colors + ) + + : m_C(C), m_N(N), m_n_isects(n_isects), m_packed(packed), + m_concat_stride(concat_stride), + m_concatenated_data(concatenated_data), m_means2d(means2d), + m_conics(conics), m_colors(colors), m_opacities(opacities), + m_backgrounds(backgrounds), m_masks(masks), + m_image_width(image_width), m_image_height(image_height), + m_tile_size(tile_size), m_tile_width(tile_width), + m_tile_height(tile_height), m_tile_offsets(tile_offsets), + m_flatten_ids(flatten_ids), m_render_alphas(render_alphas), + m_last_ids(last_ids), m_v_render_colors(v_render_colors), + m_v_render_alphas(v_render_alphas), m_v_means2d_abs(v_means2d_abs), + m_v_means2d(v_means2d), m_v_conics(v_conics), m_v_colors(v_colors), + m_v_opacities(v_opacities), m_slm_flatten_ids(slm_flatten_ids), + m_slm_means2d(slm_means2d), m_slm_opacities(slm_opacities), + m_slm_conics(slm_conics), m_slm_colors(slm_colors) {} + + [[intel::reqd_sub_group_size(16)]] + void operator()(sycl::nd_item<3> work_item) const { + // Compute camera and tile indices (each work-group corresponds to a + // tile) + const uint32_t camera_id = work_item.get_group(0); + const uint32_t tile_y = work_item.get_group(1); + const uint32_t tile_x = work_item.get_group(2); + const int32_t tile_id = tile_y * m_tile_width + tile_x; + + // Each work-work_item covers one pixel within the tile. + const uint32_t i = tile_y * m_tile_size + work_item.get_local_id(1); + const uint32_t j = tile_x * m_tile_size + work_item.get_local_id(2); + // Clamp pixel index to valid range. + const int32_t pix_id = sycl::min( + static_cast(i * m_image_width + j), + static_cast(m_image_width * m_image_height - 1) + ); + + // Adjust pointers to the current camera. + const int32_t *tile_offsets_ptr = + m_tile_offsets + camera_id * m_tile_height * m_tile_width; + + const int32_t range_start = tile_offsets_ptr[tile_id]; + int32_t range_end; + if ((camera_id == m_C - 1) && + (tile_id == static_cast(m_tile_width * m_tile_height - 1) + )) { + range_end = m_n_isects; + } else { + range_end = tile_offsets_ptr[tile_id + 1]; + } + + const S *render_alphas_ptr = + m_render_alphas + camera_id * m_image_height * m_image_width; + const int32_t *last_ids_ptr = + m_last_ids + camera_id * m_image_height * m_image_width; + const S *v_render_colors_ptr = + m_v_render_colors + + camera_id * m_image_height * m_image_width * COLOR_DIM; + const S *v_render_alphas_ptr = + m_v_render_alphas + camera_id * m_image_height * m_image_width; + const S *backgrounds_ptr = m_backgrounds; + if (backgrounds_ptr != nullptr) { + backgrounds_ptr += camera_id * COLOR_DIM; + } + const bool *masks_ptr = m_masks; + if (masks_ptr != nullptr) { + masks_ptr += camera_id * m_tile_height * m_tile_width; + } + + // If a mask exists and this tile is not active, do nothing. + if (masks_ptr != nullptr && !masks_ptr[tile_id]) { + return; + } + + // Compute the pixel’s center. + const S px = static_cast(j) + static_cast(0.5); + const S py = static_cast(i) + static_cast(0.5); + const bool inside = (i < m_image_height && j < m_image_width); + + // In the forward pass T_final = 1 - render_alphas. + const S T_final = static_cast(1.0) - render_alphas_ptr[pix_id]; + S T = T_final; + // Buffer to accumulate contributions (one per channel). + BufferType_t buffer{}; + // The index of the last gaussian that contributed (if inside). + const int32_t bin_final = inside ? last_ids_ptr[pix_id] : 0; + + // Load the pixel’s downstream gradients. + BufferType_t v_render_c; + readToBuffer(v_render_c, v_render_colors_ptr + pix_id * COLOR_DIM); + + const S v_render_a = v_render_alphas_ptr[pix_id]; + + int32_t numGaussians = range_end - range_start; + int32_t batchSize = CHUNK_SIZE; + int32_t numBatches = (numGaussians + batchSize - 1) / batchSize; + + const size_t threadRank = work_item.get_local_linear_id( + ); // given that range in 0th dimension is 1 + + for (int32_t b = numBatches - 1; b >= 0; b--) { + + work_item.barrier(sycl::access::fence_space::local_space); + + int32_t batchStart = b * batchSize + range_start; + int32_t numel = sycl::min(batchSize, range_end - batchStart); + int32_t batchEnd = batchStart + numel; + + int32_t loadIdx = batchStart + threadRank; + int32_t g_thread = -1; + if (loadIdx < range_end && threadRank < CHUNK_SIZE) { + int32_t g = m_flatten_ids[loadIdx]; + g_thread = g; + m_slm_flatten_ids[threadRank] = g; + + if constexpr (CONCAT_DATA) { + const S *data = m_concatenated_data + g * m_concat_stride; + + if constexpr (COLOR_DIM == 3) { + auto temp = + *(reinterpret_cast *>(data)); + auto temp16 = temp.template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + m_slm_means2d[threadRank] = {temp[0], temp[1]}; + m_slm_conics[threadRank] = {temp[2], temp[3], temp[4]}; + m_slm_colors[threadRank] = {temp[5], temp[6], temp[7]}; + } else { + auto xy = + *(reinterpret_cast *>(data)); + m_slm_means2d[threadRank] = xy.template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + + auto conic = *( + reinterpret_cast *>(data + 2) + ); + m_slm_conics[threadRank] = conic.template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + auto color = *(reinterpret_cast< + const BufferType_t *>( + data + 2 + 3 + )); + m_slm_colors[threadRank] = color.template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + ; + } + } + m_slm_opacities[threadRank] = + static_cast(*(data + 2 + 3 + COLOR_DIM)); + + } else { + m_slm_means2d[threadRank] = + m_means2d[g] + .template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + + m_slm_opacities[threadRank] = + static_cast(m_opacities[g]); + auto temp = *( + reinterpret_cast *>(m_conics + g) + ); + + m_slm_conics[threadRank] = temp.template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + auto temp2 = *(reinterpret_cast< + const BufferType_t *>( + m_colors + g * COLOR_DIM + )); + m_slm_colors[threadRank] = temp2.template convert< + sycl::half, + sycl::rounding_mode::automatic>(); + } + } + } + + work_item.barrier(sycl::access::fence_space::local_space); + + for (int32_t idx = numel - 1; idx >= 0; idx--) { + // Only process gaussians that actually contributed in the + // forward pass. + + bool toProcess{true}; + if (idx + batchStart > bin_final) + toProcess = false; + + const int32_t g = m_slm_flatten_ids[idx]; + + // Load forward parameters. + sycl::vec xy = + m_slm_means2d[idx] + .template convert(); + const S opac = static_cast(m_slm_opacities[idx]); + auto conic = m_slm_conics[idx] + .convert(); + + BufferType_t rgb; + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + rgb = m_slm_colors[idx] + .template convert< + S, + sycl::rounding_mode::automatic>(); + } else { + if constexpr (CONCAT_DATA) { + readToBuffer( + rgb, + m_concatenated_data + g * m_concat_stride + 2 + 3 + ); + } else { + readToBuffer(rgb, m_colors + g * COLOR_DIM); + } + } + + // Compute distance from pixel center. + sycl::vec delta = {xy.x() - px, xy.y() - py}; + S sigma = + static_cast(0.5) * (conic.x() * delta.x() * delta.x() + + conic.z() * delta.y() * delta.y()) + + conic.y() * delta.x() * delta.y(); + S vis = sycl::exp(-sigma); + S alpha = sycl::min(static_cast(0.999), opac * vis); + if (sigma < static_cast(0.0) || + alpha < static_cast(1.0 / 255.0)) + toProcess = false; + + BufferType_t v_rgb_local{}; + sycl::vec v_conic_local{}; + sycl::vec v_xy_local{}; + sycl::vec v_xy_abs_local{}; + S v_opacity_local{0.0}; + + if (toProcess) { + + // Compute reciprocal factor and update T. + const S ra = + static_cast(1.0) / (static_cast(1.0) - alpha); + T *= ra; + const S fac = alpha * T; + + // Compute gradient contribution from color. + + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + v_rgb_local[k] = fac * v_render_c[k]; + } + + // Compute partial derivative of alpha. + S v_alpha = static_cast(0.0); + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + v_alpha += + (rgb[k] * T - buffer[k] * ra) * v_render_c[k]; + } + v_alpha += T_final * ra * v_render_a; + if (backgrounds_ptr != nullptr) { + S accum = static_cast(0.0); + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + accum += backgrounds_ptr[k] * v_render_c[k]; + } + v_alpha += -T_final * ra * accum; + } + + if (opac * vis <= static_cast(0.999)) { + const S v_sigma = -opac * vis * v_alpha; + v_conic_local[0] = static_cast(0.5) * v_sigma * + delta.x() * delta.x(); + v_conic_local[1] = v_sigma * delta.x() * delta.y(); + v_conic_local[2] = static_cast(0.5) * v_sigma * + delta.y() * delta.y(); + v_xy_local[0] = v_sigma * (conic.x() * delta.x() + + conic.y() * delta.y()); + v_xy_local[1] = v_sigma * (conic.y() * delta.x() + + conic.z() * delta.y()); + if (m_v_means2d_abs != nullptr) { + v_xy_abs_local[0] = std::abs(v_xy_local[0]); + v_xy_abs_local[1] = std::abs(v_xy_local[1]); + } + v_opacity_local = vis * v_alpha; + } + + // Update the buffer. + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + buffer[k] += rgb[k] * fac; + } + } + + BufferType_t local_color; + if constexpr (BufferType::isVec) { + local_color = sycl::reduce_over_group( + work_item.get_group(), + v_rgb_local, + sycl::plus>() + ); + } else { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + local_color[k] = sycl::reduce_over_group( + work_item.get_group(), + v_rgb_local[k], + sycl::plus() + ); + } + } + + S local_opacity = sycl::reduce_over_group( + work_item.get_group(), v_opacity_local, sycl::plus() + ); + auto local_conic = sycl::reduce_over_group( + work_item.get_group(), + v_conic_local, + sycl::plus>() + ); + auto local_mean = sycl::reduce_over_group( + work_item.get_group(), + v_xy_local, + sycl::plus>() + ); + + sycl::vec local_mean_abs; + if (m_v_means2d_abs != nullptr) { + local_mean_abs = sycl::reduce_over_group( + work_item.get_group(), + v_xy_abs_local, + sycl::plus>() + ); + } + + if (threadRank == idx) { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + gpuAtomicAddGlobal( + m_v_colors[g * COLOR_DIM + k], local_color[k] + ); + } + gpuAtomicAddGlobal(m_v_opacities[g], local_opacity); + gpuAtomicAddGlobal(m_v_conics[g].x, local_conic[0]); + gpuAtomicAddGlobal(m_v_conics[g].y, local_conic[1]); + gpuAtomicAddGlobal(m_v_conics[g].z, local_conic[2]); + gpuAtomicAddGlobal(m_v_means2d[g].x(), local_mean[0]); + gpuAtomicAddGlobal(m_v_means2d[g].y(), local_mean[1]); + if (m_v_means2d_abs != nullptr) { + gpuAtomicAddGlobal( + m_v_means2d_abs[g].x(), local_mean_abs[0] + ); + gpuAtomicAddGlobal( + m_v_means2d_abs[g].y(), local_mean_abs[1] + ); + } + } + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp new file mode 100644 index 00000000..b3ae0a6a --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp @@ -0,0 +1,322 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "types.hpp" + +namespace gsplat::xpu { + +template +struct RasterizeToPixelsFwdKernel { + const uint32_t m_C; + const uint32_t m_N; + const uint32_t m_n_isects; + const bool m_packed; + const uint32_t m_concat_stride; + const S *m_concatenated_data; + const sycl::vec + *m_means2d; // [C, N, 2] or [nnz, 2] // <<< TYPE CHANGED + const vec3 *m_conics; // [C, N, 3] or [nnz, 3] // <<< TYPE CHANGED + const S *m_colors; // [C, N, COLOR_DIM] or [nnz, COLOR_DIM] + const S *m_opacities; // [C, N] or [nnz] + const S *m_backgrounds; // [C, COLOR_DIM] + const bool *m_masks; // [C, tile_height, tile_width] + const uint32_t m_image_width; + const uint32_t m_image_height; + const uint32_t m_tile_size; + const uint32_t m_tile_width; + const uint32_t m_tile_height; + const int32_t *m_tile_offsets; // [C, tile_height, tile_width] + const int32_t *m_flatten_ids; // [n_isects] + S *m_render_colors; // [C, image_height, image_width, COLOR_DIM] + S *m_render_alphas; // [C, image_height, image_width, 1] + int32_t *m_last_ids; // [C, image_height, image_width] + sycl::local_accessor m_slm_flatten_ids; + sycl::local_accessor, 1> m_slm_means2d; + sycl::local_accessor m_slm_opacities; + sycl::local_accessor, 1> m_slm_conics; + sycl::local_accessor, 1> m_slm_colors; + + RasterizeToPixelsFwdKernel( + const uint32_t C, + const uint32_t N, + const uint32_t n_isects, + const bool packed, + const uint32_t concat_stride, + const S *concatenated_data, + const sycl::vec + *means2d, // [C, N, 2] or [nnz, 2] // <<< TYPE CHANGED + const vec3 *conics, // [C, N, 3] or [nnz, 3] // <<< TYPE CHANGED + const S *colors, // [C, N, COLOR_DIM] or [nnz, COLOR_DIM] + const S *opacities, // [C, N] or [nnz] + const S *backgrounds, // [C, COLOR_DIM] + const bool *masks, // [C, tile_height, tile_width] + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + const int32_t *tile_offsets, // [C, tile_height, tile_width] + const int32_t *flatten_ids, // [n_isects] + S *render_colors, // [C, image_height, image_width, COLOR_DIM] + S *render_alphas, // [C, image_height, image_width, 1] + int32_t *last_ids, // [C, image_height, image_width] + sycl::local_accessor slm_flatten_ids, + sycl::local_accessor, 1> slm_means2d, + sycl::local_accessor slm_opacities, + sycl::local_accessor, 1> slm_conics, + sycl::local_accessor, 1> slm_colors + ) + + : m_C(C), m_N(N), m_n_isects(n_isects), m_packed(packed), + m_concat_stride(concat_stride), + m_concatenated_data(concatenated_data), m_means2d(means2d), + m_conics(conics), m_colors(colors), m_opacities(opacities), + m_backgrounds(backgrounds), m_masks(masks), + m_image_width(image_width), m_image_height(image_height), + m_tile_size(tile_size), m_tile_width(tile_width), + m_tile_height(tile_height), m_tile_offsets(tile_offsets), + m_flatten_ids(flatten_ids), m_render_colors(render_colors), + m_render_alphas(render_alphas), m_last_ids(last_ids), + m_slm_flatten_ids(slm_flatten_ids), m_slm_means2d(slm_means2d), + m_slm_opacities(slm_opacities), m_slm_conics(slm_conics), + m_slm_colors(slm_colors) {} + + [[intel::reqd_sub_group_size(16)]] + void operator()(sycl::nd_item<3> work_item) const { + + const uint32_t camera_id = work_item.get_group(0); // [0, C) + const uint32_t tile_y = work_item.get_group(1); // [0, tile_height) + const uint32_t tile_x = work_item.get_group(2); // [0, tile_width) + const int32_t tile_id = tile_y * m_tile_width + tile_x; + + const int32_t *tile_offsets_ptr = + m_tile_offsets + camera_id * m_tile_height * m_tile_width; + + const int32_t range_start = tile_offsets_ptr[tile_id]; + int32_t range_end = 0; + + if ((camera_id == m_C - 1) && + (tile_id == static_cast(m_tile_width * m_tile_height - 1) + )) { + range_end = m_n_isects; + } else { + range_end = tile_offsets_ptr[tile_id + 1]; + } + + S *render_colors_ptr = m_render_colors + camera_id * m_image_height * + m_image_width * COLOR_DIM; + S *render_alphas_ptr = + m_render_alphas + camera_id * m_image_height * m_image_width; + int32_t *last_ids_ptr = + m_last_ids + camera_id * m_image_height * m_image_width; + + BufferType_t backgroundColor{}; + if (m_backgrounds != nullptr) { + readToBuffer( + backgroundColor, m_backgrounds + camera_id * COLOR_DIM + ); + } + const bool *masks_ptr = m_masks; + if (masks_ptr != nullptr) { + masks_ptr += camera_id * m_tile_height * m_tile_width; + } + + // Local range is {1, tile_size, tile_size} so that: + // local_id(1) in [0, tile_size), local_id(2) in [0, tile_size) + const uint32_t i = tile_y * m_tile_size + work_item.get_local_id(1); + const uint32_t j = tile_x * m_tile_size + work_item.get_local_id(2); + const int32_t pix_id = i * m_image_width + j; + // Compute pixel center + bool inside = (i < m_image_height && j < m_image_width); + bool done = !inside; + + // If a mask exists and the tile is marked false, output background + // color immediately. + if (masks_ptr != nullptr && inside && !masks_ptr[tile_id]) { + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + render_colors_ptr[pix_id * COLOR_DIM + k] = backgroundColor[k]; + } + return; + } + + // Initialize transmittance and pixel accumulator. + S T = static_cast(1.0); + + BufferType_t pix_out{}; + + int32_t cur_idx = 0; + + int32_t numGaussians = range_end - range_start; + int32_t batchSize = CHUNK_SIZE; + int32_t numBatches = (numGaussians + batchSize - 1) / batchSize; + + const size_t localId_y = work_item.get_local_id(1); + const size_t localId_x = work_item.get_local_id(2); + const size_t groupWidth = work_item.get_local_range(2); + const size_t threadRank = + localId_y * groupWidth + + localId_x; // given that range in 0th dimension is 1 + + // Compute pixel coordinates: each work-item covers one pixel inside the + // tile. + const S px = static_cast(j) + static_cast(0.5); + const S py = static_cast(i) + static_cast(0.5); + + for (uint32_t b = 0; b < numBatches; b++) { + + work_item.barrier(sycl::access::fence_space::local_space); + + int32_t batchStart = b * batchSize + range_start; + int32_t idx = batchStart + threadRank; + + if (idx < range_end && threadRank < CHUNK_SIZE) { + + int32_t g = m_flatten_ids[idx]; + m_slm_flatten_ids[threadRank] = g; + + if constexpr (CONCAT_DATA) { + const S *data = m_concatenated_data + g * m_concat_stride; + if constexpr (COLOR_DIM == 3) { + // means(2) + conics(3) + colors(3) + opac(1) + const S *data = + m_concatenated_data + g * m_concat_stride; + auto temp = + *(reinterpret_cast *>(data)); + + m_slm_means2d[threadRank] = {temp[0], temp[1]}; + m_slm_conics[threadRank] = {temp[2], temp[3], temp[4]}; + m_slm_colors[threadRank] = {temp[5], temp[6], temp[7]}; + m_slm_opacities[threadRank] = + *(data + 2 + 3 + COLOR_DIM); + } else { + if constexpr (BufferType::isVec && + COLOR_DIM == 4) { + // means(2) + conics(3) + colors(4) + opac(1) + auto temp1 = + *(reinterpret_cast *>(data + )); + auto temp2 = + *(reinterpret_cast *>( + data + 8 + )); + m_slm_means2d[threadRank] = {temp1[0], temp1[1]}; + m_slm_conics[threadRank] = { + temp1[2], temp1[3], temp1[4] + }; + m_slm_colors[threadRank] = { + temp1[5], temp1[6], temp1[7], temp2[0] + }; + m_slm_opacities[threadRank] = temp2[1]; + + } else { + m_slm_means2d[threadRank] = + *(reinterpret_cast *>(data + )); + m_slm_conics[threadRank] = + *(reinterpret_cast *>( + data + 2 + )); + m_slm_colors[threadRank] = + *(reinterpret_cast< + const BufferType_t *>( + data + 2 + 3 + )); + m_slm_opacities[threadRank] = + *(data + 2 + 3 + COLOR_DIM); + } + } + + } else { + m_slm_means2d[threadRank] = m_means2d[g]; + m_slm_opacities[threadRank] = m_opacities[g]; + m_slm_conics[threadRank] = *( + reinterpret_cast *>(m_conics + g) + ); + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + m_slm_colors[threadRank] = + *(reinterpret_cast + *>(m_colors + g * COLOR_DIM) + ); + } + } + } + + work_item.barrier(sycl::access::fence_space::local_space); + + int32_t rangeDiff = range_end - batchStart; + int32_t endSize = (rangeDiff < batchSize) ? rangeDiff : batchSize; + + for (int i = 0; i < endSize && (!done); i++) { + + int32_t g = m_slm_flatten_ids[i]; + const sycl::vec xy = m_slm_means2d[i]; + const S opac = m_slm_opacities[i]; + const auto conic = m_slm_conics[i]; + + sycl::vec delta = {xy[0] - px, xy[1] - py}; + S sigma = + static_cast(0.5) * (conic.x() * delta.x() * delta.x() + + conic.z() * delta.y() * delta.y()) + + conic.y() * delta.x() * delta.y(); + + S alpha = + sycl::min(static_cast(0.999), opac * sycl::exp(-sigma)); + + if (sigma < static_cast(0.0) || + alpha < static_cast(1.0 / 255.0)) + continue; + + S next_T = T * (static_cast(1.0) - alpha); + if (next_T <= static_cast(1e-4)) { + done = true; + break; + } + + const S vis = alpha * T; + + BufferType_t currColor; + if constexpr (BufferType::isVec && + COLOR_DIM <= 4) { + currColor = m_slm_colors[i]; + } else { + if constexpr (CONCAT_DATA) { + readToBuffer( + currColor, + m_concatenated_data + g * m_concat_stride + 2 + 3 + ); + } else { + readToBuffer(currColor, m_colors + g * COLOR_DIM); + } + } + + pix_out += currColor * vis; + + cur_idx = batchStart + i; + T = next_T; + } + } + + // Write out results if the pixel is within the image. + if (inside) { + + render_alphas_ptr[pix_id] = static_cast(1.0) - T; + last_ids_ptr[pix_id] = cur_idx; + + S *current_pixel_color_ptr_base = + render_colors_ptr + pix_id * COLOR_DIM; + auto *current_pixel_color_ptr = + reinterpret_cast *>( + current_pixel_color_ptr_base + ); + +#pragma unroll + for (uint32_t k = 0; k < COLOR_DIM; ++k) { + current_pixel_color_ptr[0][k] = + pix_out[k] + T * backgroundColor[k]; + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/RelocationKernel.hpp b/gsplat/sycl/include/kernels/RelocationKernel.hpp new file mode 100644 index 00000000..9fc88eb0 --- /dev/null +++ b/gsplat/sycl/include/kernels/RelocationKernel.hpp @@ -0,0 +1,57 @@ +#pragma once +#include + +namespace gsplat::xpu::kernels { + +template class RelocationKernel { + private: + const scalar_t *opacities; + const scalar_t *scales; + const int *ratios; + const scalar_t *binoms; + const int n_max; + scalar_t *new_opacities; + scalar_t *new_scales; + + public: + RelocationKernel( + const scalar_t *opacities, + const scalar_t *scales, + const int *ratios, + const scalar_t *binoms, + const int n_max, + scalar_t *new_opacities, + scalar_t *new_scales + ) + : opacities(opacities), scales(scales), ratios(ratios), binoms(binoms), + n_max(n_max), new_opacities(new_opacities), new_scales(new_scales) {} + + void operator()(sycl::id<1> item) const { + int idx = item[0]; + + int n_idx = ratios[idx]; + float denom_sum = 0.0f; + + // compute new opacity + new_opacities[idx] = + 1.0f - + sycl::pow(1.0f - static_cast(opacities[idx]), 1.0f / n_idx); + + // compute new scale + for (int i = 1; i <= n_idx; ++i) { + for (int k = 0; k <= (i - 1); ++k) { + float bin_coeff = binoms[(i - 1) * n_max + k]; + float term = + (sycl::pow(-1.0f, k) / sycl::sqrt(static_cast(k + 1)) + ) * + sycl::pow(static_cast(new_opacities[idx]), k + 1); + denom_sum += (bin_coeff * term); + } + } + float coeff = (opacities[idx] / denom_sum); + for (int i = 0; i < 3; ++i) + new_scales[idx * 3 + i] = coeff * scales[idx * 3 + i]; + } +}; + +} // namespace gsplat::xpu::kernels diff --git a/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp new file mode 100644 index 00000000..0d8c4584 --- /dev/null +++ b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp @@ -0,0 +1,120 @@ +#pragma once + +#include "Sycl_utils.hpp" +#include "transform.hpp" +#include "types.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { + +template struct WorldToCamBwdKernel { + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [N, 3] + const T *m_covars; // [N, 3, 3] + const T *m_viewmats; // [C, 4, 4] + const T *m_v_means_c; // [C, N, 3] + const T *m_v_covars_c; // [C, N, 3, 3] + T *m_v_means; // [N, 3] + T *m_v_covars; // [N, 3, 3] + T *m_v_viewmats; // [C, 4, 4] + + WorldToCamBwdKernel( + const uint32_t C, + const uint32_t N, + const T *means, + const T *covars, + const T *viewmats, + const T *v_means_c, + const T *v_covars_c, + T *v_means, + T *v_covars, + T *v_viewmats + ) + : m_C(C), m_N(N), m_means(means), m_covars(covars), + m_viewmats(viewmats), m_v_means_c(v_means_c), + m_v_covars_c(v_covars_c), m_v_means(v_means), m_v_covars(v_covars), + m_v_viewmats(v_viewmats) {} + + void operator()(sycl::nd_item<1> work_item) const { + const uint32_t idx = work_item.get_global_id(0); + + if (idx >= m_C * m_N) { + return; + } + + const uint32_t cid = idx / m_N; // camera id + const uint32_t gid = idx % m_N; // gaussian id + + // shift pointers to the current camera and gaussian + const T *means = m_means + (gid * 3); + const T *covars = m_covars + (gid * 9); + const T *viewmats = m_viewmats + (cid * 16); + + // glm is column-major but input is row-major + const mat3 R = mat3( + viewmats[0], + viewmats[4], + viewmats[8], // 1st column + viewmats[1], + viewmats[5], + viewmats[9], // 2nd column + viewmats[2], + viewmats[6], + viewmats[10] // 3rd column + ); + + const vec3 t = vec3(viewmats[3], viewmats[7], viewmats[11]); + + vec3 v_mean(0.f); + mat3 v_covar(0.f); + mat3 v_R(0.f); + vec3 v_t(0.f); + + if (m_v_means_c != nullptr) { + const vec3 v_mean_c = glm::make_vec3(m_v_means_c + (idx * 3)); + const vec3 mean = glm::make_vec3(means); + pos_world_to_cam_vjp(R, t, mean, v_mean_c, v_R, v_t, v_mean); + } + if (m_v_covars_c != nullptr) { + const mat3 v_covar_c_t = + glm::make_mat3(m_v_covars_c + (idx * 9)); + const mat3 v_covar_c = glm::transpose(v_covar_c_t); + const mat3 covar = glm::make_mat3(covars); + covar_world_to_cam_vjp(R, covar, v_covar_c, v_R, v_covar); + } + + if (m_v_means != nullptr) { + T *v_means = m_v_means + (gid * 3); +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { + gpuAtomicAdd(v_means + i, v_mean[i]); + } + } + + if (m_v_covars != nullptr) { + T *v_covars = m_v_covars + (gid * 9); +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + gpuAtomicAdd(v_covars + i * 3 + j, v_covar[j][i]); + } + } + } + + if (m_v_viewmats != nullptr) { + T *v_viewmats = m_v_viewmats + cid * 16; +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + gpuAtomicAdd(v_viewmats + i * 4 + j, v_R[j][i]); + } + gpuAtomicAdd(v_viewmats + i * 4 + 3, v_t[i]); + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp b/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp new file mode 100644 index 00000000..d272f63d --- /dev/null +++ b/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp @@ -0,0 +1,92 @@ +#pragma once + +/**************************************************************************** + * World to Camera Transformation Forward Pass + * From: + *https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/csrc/world_to_cam_fwd.cu + ****************************************************************************/ + +#include "transform.hpp" +#include "types.hpp" + +namespace gsplat::xpu { + +template struct WorldToCamFwdKernel { + const uint32_t m_C; + const uint32_t m_N; + const T *m_means; // [N, 3] + const T *m_covars; // [N, 3, 3] + const T *m_viewmats; // [C, 4, 4] + T *m_means_c; // [C, N, 3] + T *m_covars_c; // [C, N, 3, 3] + + WorldToCamFwdKernel( + const uint32_t C, + const uint32_t N, + const T *means, + const T *covars, + const T *viewmats, + T *means_c, + T *covars_c + ) + : m_C(C), m_N(N), m_means(means), m_covars(covars), + m_viewmats(viewmats), m_means_c(means_c), m_covars_c(covars_c) {} + + void operator()(sycl::nd_item<1> work_item) const { + const int64_t idx = work_item.get_global_id(0); + if (idx >= m_C * m_N) { + return; + } + + const uint32_t cid = idx / m_N; // camera id + const uint32_t gid = idx % m_N; // gaussian id + + // shift pointers to the current camera and gaussian + const T *means = m_means + (gid * 3); + const T *covars = m_covars + (gid * 9); + const T *viewmats = m_viewmats + (cid * 16); + + // glm is column-major but input is row-major + const mat3 R = mat3( + viewmats[0], + viewmats[4], + viewmats[8], // 1st column + viewmats[1], + viewmats[5], + viewmats[9], // 2nd column + viewmats[2], + viewmats[6], + viewmats[10] // 3rd column + ); + + const vec3 t = vec3(viewmats[3], viewmats[7], viewmats[11]); + + if (m_means_c != nullptr) { + vec3 mean_c; + const vec3 mean = glm::make_vec3(means); + pos_world_to_cam(R, t, mean, mean_c); + T *means_c = m_means_c + (idx * 3); +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows + means_c[i] = mean_c[i]; + } + } + + // write to outputs: glm is column-major but we want row-major + if (m_covars_c != nullptr) { + mat3 covar_c; + const mat3 covar = glm::make_mat3(covars); + covar_world_to_cam(R, covar, covar_c); + T *covars_c = m_covars_c + (idx * 9); +#pragma unroll + for (uint32_t i = 0; i < 3; i++) { // rows +#pragma unroll + for (uint32_t j = 0; j < 3; j++) { // cols + covars_c[i * 3 + j] = T(covar_c[j][i]); + } + } + } + } +}; + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/proj.hpp b/gsplat/sycl/include/proj.hpp new file mode 100644 index 00000000..4ed9890f --- /dev/null +++ b/gsplat/sycl/include/proj.hpp @@ -0,0 +1,340 @@ +#pragma once + +#include "types.hpp" + +template +inline void ortho_proj( + // inputs + const vec3 mean3d, + const mat3 cov3d, + const T fx, + const T fy, + const T cx, + const T cy, + const uint32_t width, + const uint32_t height, + // outputs + mat2 &cov2d, + vec2 &mean2d +) { + T x = mean3d[0], y = mean3d[1]; // z = mean3d[2]; + + // mat3x2 is 3 columns x 2 rows. + mat3x2 J = mat3x2( + fx, + 0.f, // 1st column + 0.f, + fy, // 2nd column + 0.f, + 0.f // 3rd column + ); + cov2d = J * cov3d * glm::transpose(J); + mean2d = vec2({fx * x + cx, fy * y + cy}); +} + +template +inline void ortho_proj_vjp( + // fwd inputs + const vec3 mean3d, + const mat3 cov3d, + const T fx, + const T fy, + const T cx, + const T cy, + const uint32_t width, + const uint32_t height, + // grad outputs + const mat2 v_cov2d, + const vec2 v_mean2d, + // grad inputs + vec3 &v_mean3d, + mat3 &v_cov3d +) { + // T x = mean3d[0], y = mean3d[1], z = mean3d[2]; + + // mat3x2 is 3 columns x 2 rows. + mat3x2 J = mat3x2( + fx, + 0.f, // 1st column + 0.f, + fy, // 2nd column + 0.f, + 0.f // 3rd column + ); + + // cov = J * V * Jt; G = df/dcov = v_cov + // -> df/dV = Jt * G * J + // -> df/dJ = G * J * Vt + Gt * J * V + v_cov3d += glm::transpose(J) * v_cov2d * J; + + // df/dx = fx * df/dpixx + // df/dy = fy * df/dpixy + // df/dz = 0 + v_mean3d += vec3(fx * v_mean2d[0], fy * v_mean2d[1], 0.f); +} + +template +inline void persp_proj( + // inputs + const vec3 mean3d, + const mat3 cov3d, + const T fx, + const T fy, + const T cx, + const T cy, + const uint32_t width, + const uint32_t height, + // outputs + mat2 &cov2d, + vec2 &mean2d +) { + T x = mean3d[0], y = mean3d[1], z = mean3d[2]; + + T tan_fovx = 0.5f * width / fx; + T tan_fovy = 0.5f * height / fy; + T lim_x_pos = (width - cx) / fx + 0.3f * tan_fovx; + T lim_x_neg = cx / fx + 0.3f * tan_fovx; + T lim_y_pos = (height - cy) / fy + 0.3f * tan_fovy; + T lim_y_neg = cy / fy + 0.3f * tan_fovy; + + T rz = 1.f / z; + T rz2 = rz * rz; + T tx = z * sycl::min(lim_x_pos, sycl::max(-lim_x_neg, x * rz)); + T ty = z * sycl::min(lim_y_pos, sycl::max(-lim_y_neg, y * rz)); + + // mat3x2 is 3 columns x 2 rows. + mat3x2 J = mat3x2( + fx * rz, + 0.f, // 1st column + 0.f, + fy * rz, // 2nd column + -fx * tx * rz2, + -fy * ty * rz2 // 3rd column + ); + cov2d = J * cov3d * glm::transpose(J); + mean2d = vec2({fx * x * rz + cx, fy * y * rz + cy}); +} + +template +inline void persp_proj_vjp( + // fwd inputs + const vec3 mean3d, + const mat3 cov3d, + const T fx, + const T fy, + const T cx, + const T cy, + const uint32_t width, + const uint32_t height, + // grad outputs + const mat2 v_cov2d, + const vec2 v_mean2d, + // grad inputs + vec3 &v_mean3d, + mat3 &v_cov3d +) { + T x = mean3d[0], y = mean3d[1], z = mean3d[2]; + + T tan_fovx = 0.5f * width / fx; + T tan_fovy = 0.5f * height / fy; + T lim_x_pos = (width - cx) / fx + 0.3f * tan_fovx; + T lim_x_neg = cx / fx + 0.3f * tan_fovx; + T lim_y_pos = (height - cy) / fy + 0.3f * tan_fovy; + T lim_y_neg = cy / fy + 0.3f * tan_fovy; + + T rz = 1.f / z; + T rz2 = rz * rz; + T tx = z * sycl::min(lim_x_pos, sycl::max(-lim_x_neg, x * rz)); + T ty = z * sycl::min(lim_y_pos, sycl::max(-lim_y_neg, y * rz)); + + // mat3x2 is 3 columns x 2 rows. + mat3x2 J = mat3x2( + fx * rz, + 0.f, // 1st column + 0.f, + fy * rz, // 2nd column + -fx * tx * rz2, + -fy * ty * rz2 // 3rd column + ); + + // cov = J * V * Jt; G = df/dcov = v_cov + // -> df/dV = Jt * G * J + // -> df/dJ = G * J * Vt + Gt * J * V + v_cov3d += glm::transpose(J) * v_cov2d * J; + + // df/dx = fx * rz * df/dpixx + // df/dy = fy * rz * df/dpixy + // df/dz = - fx * mean.x * rz2 * df/dpixx - fy * mean.y * rz2 * df/dpixy + v_mean3d += vec3( + fx * rz * v_mean2d[0], + fy * rz * v_mean2d[1], + -(fx * x * v_mean2d[0] + fy * y * v_mean2d[1]) * rz2 + ); + + // df/dx = -fx * rz2 * df/dJ_02 + // df/dy = -fy * rz2 * df/dJ_12 + // df/dz = -fx * rz2 * df/dJ_00 - fy * rz2 * df/dJ_11 + // + 2 * fx * tx * rz3 * df/dJ_02 + 2 * fy * ty * rz3 + T rz3 = rz2 * rz; + mat3x2 v_J = v_cov2d * J * glm::transpose(cov3d) + + glm::transpose(v_cov2d) * J * cov3d; + + // fov clipping + if (x * rz <= lim_x_pos && x * rz >= -lim_x_neg) { + v_mean3d.x += -fx * rz2 * v_J[2][0]; + } else { + v_mean3d.z += -fx * rz3 * v_J[2][0] * tx; + } + if (y * rz <= lim_y_pos && y * rz >= -lim_y_neg) { + v_mean3d.y += -fy * rz2 * v_J[2][1]; + } else { + v_mean3d.z += -fy * rz3 * v_J[2][1] * ty; + } + v_mean3d.z += -fx * rz2 * v_J[0][0] - fy * rz2 * v_J[1][1] + + 2.f * fx * tx * rz3 * v_J[2][0] + + 2.f * fy * ty * rz3 * v_J[2][1]; +} + +template +inline void fisheye_proj( + // inputs + const vec3 mean3d, + const mat3 cov3d, + const T fx, + const T fy, + const T cx, + const T cy, + const uint32_t width, + const uint32_t height, + // outputs + mat2 &cov2d, + vec2 &mean2d +) { + T x = mean3d[0], y = mean3d[1], z = mean3d[2]; + + T eps = 0.0000001f; + T xy_len = glm::length(glm::vec2({x, y})) + eps; + T theta = glm::atan(xy_len, z + eps); + mean2d = + vec2({x * fx * theta / xy_len + cx, y * fy * theta / xy_len + cy}); + + T x2 = x * x + eps; + T y2 = y * y; + T xy = x * y; + T x2y2 = x2 + y2; + T x2y2z2_inv = 1.f / (x2y2 + z * z); + + T b = glm::atan(xy_len, z) / xy_len / x2y2; + T a = z * x2y2z2_inv / (x2y2); + mat3x2 J = mat3x2( + fx * (x2 * a + y2 * b), + fy * xy * (a - b), + fx * xy * (a - b), + fy * (y2 * a + x2 * b), + -fx * x * x2y2z2_inv, + -fy * y * x2y2z2_inv + ); + cov2d = J * cov3d * glm::transpose(J); +} + +template +inline void fisheye_proj_vjp( + // fwd inputs + const vec3 mean3d, + const mat3 cov3d, + const T fx, + const T fy, + const T cx, + const T cy, + const uint32_t width, + const uint32_t height, + // grad outputs + const mat2 v_cov2d, + const vec2 v_mean2d, + // grad inputs + vec3 &v_mean3d, + mat3 &v_cov3d +) { + T x = mean3d[0], y = mean3d[1], z = mean3d[2]; + + const T eps = 0.0000001f; + T x2 = x * x + eps; + T y2 = y * y; + T xy = x * y; + T x2y2 = x2 + y2; + T len_xy = length(glm::vec2({x, y})) + eps; + const T x2y2z2 = x2y2 + z * z; + T x2y2z2_inv = 1.f / x2y2z2; + T b = glm::atan(len_xy, z) / len_xy / x2y2; + T a = z * x2y2z2_inv / (x2y2); + v_mean3d += vec3( + fx * (x2 * a + y2 * b) * v_mean2d[0] + fy * xy * (a - b) * v_mean2d[1], + fx * xy * (a - b) * v_mean2d[0] + fy * (y2 * a + x2 * b) * v_mean2d[1], + -fx * x * x2y2z2_inv * v_mean2d[0] - fy * y * x2y2z2_inv * v_mean2d[1] + ); + + const T theta = glm::atan(len_xy, z); + const T J_b = theta / len_xy / x2y2; + const T J_a = z * x2y2z2_inv / (x2y2); + // mat3x2 is 3 columns x 2 rows. + mat3x2 J = mat3x2( + fx * (x2 * J_a + y2 * J_b), + fy * xy * (J_a - J_b), // 1st column + fx * xy * (J_a - J_b), + fy * (y2 * J_a + x2 * J_b), // 2nd column + -fx * x * x2y2z2_inv, + -fy * y * x2y2z2_inv // 3rd column + ); + v_cov3d += glm::transpose(J) * v_cov2d * J; + + mat3x2 v_J = v_cov2d * J * glm::transpose(cov3d) + + glm::transpose(v_cov2d) * J * cov3d; + T l4 = x2y2z2 * x2y2z2; + + T E = -l4 * x2y2 * theta + x2y2z2 * x2y2 * len_xy * z; + T F = 3 * l4 * theta - 3 * x2y2z2 * len_xy * z - 2 * x2y2 * len_xy * z; + + T A = x * (3 * E + x2 * F); + T B = y * (E + x2 * F); + T C = x * (E + y2 * F); + T D = y * (3 * E + y2 * F); + + T S1 = x2 - y2 - z * z; + T S2 = y2 - x2 - z * z; + T inv1 = x2y2z2_inv * x2y2z2_inv; + T inv2 = inv1 / (x2y2 * x2y2 * len_xy); + + T dJ_dx00 = fx * A * inv2; + T dJ_dx01 = fx * B * inv2; + T dJ_dx02 = fx * S1 * inv1; + T dJ_dx10 = fy * B * inv2; + T dJ_dx11 = fy * C * inv2; + T dJ_dx12 = 2.f * fy * xy * inv1; + + T dJ_dy00 = dJ_dx01; + T dJ_dy01 = fx * C * inv2; + T dJ_dy02 = 2.f * fx * xy * inv1; + T dJ_dy10 = dJ_dx11; + T dJ_dy11 = fy * D * inv2; + T dJ_dy12 = fy * S2 * inv1; + + T dJ_dz00 = dJ_dx02; + T dJ_dz01 = dJ_dy02; + T dJ_dz02 = 2.f * fx * x * z * inv1; + T dJ_dz10 = dJ_dx12; + T dJ_dz11 = dJ_dy12; + T dJ_dz12 = 2.f * fy * y * z * inv1; + + T dL_dtx_raw = dJ_dx00 * v_J[0][0] + dJ_dx01 * v_J[1][0] + + dJ_dx02 * v_J[2][0] + dJ_dx10 * v_J[0][1] + + dJ_dx11 * v_J[1][1] + dJ_dx12 * v_J[2][1]; + T dL_dty_raw = dJ_dy00 * v_J[0][0] + dJ_dy01 * v_J[1][0] + + dJ_dy02 * v_J[2][0] + dJ_dy10 * v_J[0][1] + + dJ_dy11 * v_J[1][1] + dJ_dy12 * v_J[2][1]; + T dL_dtz_raw = dJ_dz00 * v_J[0][0] + dJ_dz01 * v_J[1][0] + + dJ_dz02 * v_J[2][0] + dJ_dz10 * v_J[0][1] + + dJ_dz11 * v_J[1][1] + dJ_dz12 * v_J[2][1]; + v_mean3d.x += dL_dtx_raw; + v_mean3d.y += dL_dty_raw; + v_mean3d.z += dL_dtz_raw; +} diff --git a/gsplat/sycl/include/quat.hpp b/gsplat/sycl/include/quat.hpp new file mode 100644 index 00000000..9282bd81 --- /dev/null +++ b/gsplat/sycl/include/quat.hpp @@ -0,0 +1,53 @@ +#pragma once + +#include "types.hpp" + +template inline mat3 quat_to_rotmat(const vec4 quat) { + T w = quat[0], x = quat[1], y = quat[2], z = quat[3]; + // normalize + T inv_norm = sycl::rsqrt(x * x + y * y + z * z + w * w); + x *= inv_norm; + y *= inv_norm; + z *= inv_norm; + w *= inv_norm; + T x2 = x * x, y2 = y * y, z2 = z * z; + T xy = x * y, xz = x * z, yz = y * z; + T wx = w * x, wy = w * y, wz = w * z; + return mat3( + (1.f - 2.f * (y2 + z2)), + (2.f * (xy + wz)), + (2.f * (xz - wy)), // 1st col + (2.f * (xy - wz)), + (1.f - 2.f * (x2 + z2)), + (2.f * (yz + wx)), // 2nd col + (2.f * (xz + wy)), + (2.f * (yz - wx)), + (1.f - 2.f * (x2 + y2)) // 3rd col + ); +} + +template +inline void +quat_to_rotmat_vjp(const vec4 quat, const mat3 v_R, vec4 &v_quat) { + T w = quat[0], x = quat[1], y = quat[2], z = quat[3]; + // normalize + T inv_norm = sycl::rsqrt(x * x + y * y + z * z + w * w); + x *= inv_norm; + y *= inv_norm; + z *= inv_norm; + w *= inv_norm; + vec4 v_quat_n = vec4( + 2.f * (x * (v_R[1][2] - v_R[2][1]) + y * (v_R[2][0] - v_R[0][2]) + + z * (v_R[0][1] - v_R[1][0])), + 2.f * + (-2.f * x * (v_R[1][1] + v_R[2][2]) + y * (v_R[0][1] + v_R[1][0]) + + z * (v_R[0][2] + v_R[2][0]) + w * (v_R[1][2] - v_R[2][1])), + 2.f * (x * (v_R[0][1] + v_R[1][0]) - 2.f * y * (v_R[0][0] + v_R[2][2]) + + z * (v_R[1][2] + v_R[2][1]) + w * (v_R[2][0] - v_R[0][2])), + 2.f * (x * (v_R[0][2] + v_R[2][0]) + y * (v_R[1][2] + v_R[2][1]) - + 2.f * z * (v_R[0][0] + v_R[1][1]) + w * (v_R[0][1] - v_R[1][0])) + ); + + vec4 quat_n = vec4(w, x, y, z); + v_quat += (v_quat_n - glm::dot(v_quat_n, quat_n) * quat_n) * inv_norm; +} \ No newline at end of file diff --git a/gsplat/sycl/include/quat_scale_to_covar_preci.hpp b/gsplat/sycl/include/quat_scale_to_covar_preci.hpp new file mode 100644 index 00000000..42ddb7bd --- /dev/null +++ b/gsplat/sycl/include/quat_scale_to_covar_preci.hpp @@ -0,0 +1,121 @@ +#pragma once + +#include "quat.hpp" +#include "types.hpp" + +template +inline void quat_scale_to_covar_preci( + const vec4 quat, + const vec3 scale, + // optional outputs + mat3 *covar, + mat3 *preci +) { + mat3 R = quat_to_rotmat(quat); + if (covar != nullptr) { + // C = R * S * S * Rt + mat3 S = + mat3(scale[0], 0.f, 0.f, 0.f, scale[1], 0.f, 0.f, 0.f, scale[2]); + mat3 M = R * S; + *covar = M * glm::transpose(M); + } + if (preci != nullptr) { + // P = R * S^-1 * S^-1 * Rt + mat3 S = mat3( + 1.0f / scale[0], + 0.f, + 0.f, + 0.f, + 1.0f / scale[1], + 0.f, + 0.f, + 0.f, + 1.0f / scale[2] + ); + mat3 M = R * S; + *preci = M * glm::transpose(M); + } +} + +template +inline void quat_scale_to_covar_vjp( + // fwd inputs + const vec4 quat, + const vec3 scale, + // precompute + const mat3 R, + // grad outputs + const mat3 v_covar, + // grad inputs + vec4 &v_quat, + vec3 &v_scale +) { + // T w = quat[0], x = quat[1], y = quat[2], z = quat[3]; + T sx = scale[0], sy = scale[1], sz = scale[2]; + + // M = R * S + mat3 S = mat3(sx, 0.f, 0.f, 0.f, sy, 0.f, 0.f, 0.f, sz); + mat3 M = R * S; + + // https://math.stackexchange.com/a/3850121 + // for D = W * X, G = df/dD + // df/dW = G * XT, df/dX = WT * G + // so + // for D = M * Mt, + // df/dM = df/dM + df/dMt = G * M + (Mt * G)t = G * M + Gt * M + mat3 v_M = (v_covar + glm::transpose(v_covar)) * M; + mat3 v_R = v_M * S; + + // grad for (quat, scale) from covar + quat_to_rotmat_vjp(quat, v_R, v_quat); + + v_scale[0] += + R[0][0] * v_M[0][0] + R[0][1] * v_M[0][1] + R[0][2] * v_M[0][2]; + v_scale[1] += + R[1][0] * v_M[1][0] + R[1][1] * v_M[1][1] + R[1][2] * v_M[1][2]; + v_scale[2] += + R[2][0] * v_M[2][0] + R[2][1] * v_M[2][1] + R[2][2] * v_M[2][2]; +} + +template +inline void quat_scale_to_preci_vjp( + // fwd inputs + const vec4 quat, + const vec3 scale, + // precompute + const mat3 R, + // grad outputs + const mat3 v_preci, + // grad inputs + vec4 &v_quat, + vec3 &v_scale +) { + // T w = quat[0], x = quat[1], y = quat[2], z = quat[3]; + T sx = 1.0f / scale[0], sy = 1.0f / scale[1], sz = 1.0f / scale[2]; + + // M = R * S + mat3 S = mat3(sx, 0.f, 0.f, 0.f, sy, 0.f, 0.f, 0.f, sz); + mat3 M = R * S; + + // https://math.stackexchange.com/a/3850121 + // for D = W * X, G = df/dD + // df/dW = G * XT, df/dX = WT * G + // so + // for D = M * Mt, + // df/dM = df/dM + df/dMt = G * M + (Mt * G)t = G * M + Gt * M + mat3 v_M = (v_preci + glm::transpose(v_preci)) * M; + mat3 v_R = v_M * S; + + // grad for (quat, scale) from preci + quat_to_rotmat_vjp(quat, v_R, v_quat); + + v_scale[0] += + -sx * sx * + (R[0][0] * v_M[0][0] + R[0][1] * v_M[0][1] + R[0][2] * v_M[0][2]); + v_scale[1] += + -sy * sy * + (R[1][0] * v_M[1][0] + R[1][1] * v_M[1][1] + R[1][2] * v_M[1][2]); + v_scale[2] += + -sz * sz * + (R[2][0] * v_M[2][0] + R[2][1] * v_M[2][1] + R[2][2] * v_M[2][2]); +} diff --git a/gsplat/sycl/include/spherical_harmonics.hpp b/gsplat/sycl/include/spherical_harmonics.hpp new file mode 100644 index 00000000..9efe6207 --- /dev/null +++ b/gsplat/sycl/include/spherical_harmonics.hpp @@ -0,0 +1,357 @@ +#pragma once + +#include "types.hpp" + +// Evaluate spherical harmonics bases at unit direction for high orders using +// approach described by Efficient Spherical Harmonic Evaluation, Peter-Pike +// Sloan, JCGT 2013 See https://jcgt.org/published/0002/02/06/ for reference +// implementation +template +inline void sh_coeffs_to_color_fast( + const uint32_t degree, // degree of SH to be evaluated + const uint32_t c, // color channel + const vec3 &dir, // [3] + const T *coeffs, // [K, 3] + // output + T *colors // [3] +) { + T result = 0.2820947917738781f * coeffs[c]; + if (degree >= 1) { + T inorm = sycl::rsqrt(dir.x * dir.x + dir.y * dir.y + dir.z * dir.z); + T x = dir.x * inorm; + T y = dir.y * inorm; + T z = dir.z * inorm; + + result += + 0.48860251190292f * (-y * coeffs[1 * 3 + c] + + z * coeffs[2 * 3 + c] - x * coeffs[3 * 3 + c]); + if (degree >= 2) { + T z2 = z * z; + + T fTmp0B = -1.092548430592079f * z; + T fC1 = x * x - y * y; + T fS1 = 2.f * x * y; + T pSH6 = (0.9461746957575601f * z2 - 0.3153915652525201f); + T pSH7 = fTmp0B * x; + T pSH5 = fTmp0B * y; + T pSH8 = 0.5462742152960395f * fC1; + T pSH4 = 0.5462742152960395f * fS1; + + result += pSH4 * coeffs[4 * 3 + c] + pSH5 * coeffs[5 * 3 + c] + + pSH6 * coeffs[6 * 3 + c] + pSH7 * coeffs[7 * 3 + c] + + pSH8 * coeffs[8 * 3 + c]; + if (degree >= 3) { + T fTmp0C = -2.285228997322329f * z2 + 0.4570457994644658f; + T fTmp1B = 1.445305721320277f * z; + T fC2 = x * fC1 - y * fS1; + T fS2 = x * fS1 + y * fC1; + T pSH12 = z * (1.865881662950577f * z2 - 1.119528997770346f); + T pSH13 = fTmp0C * x; + T pSH11 = fTmp0C * y; + T pSH14 = fTmp1B * fC1; + T pSH10 = fTmp1B * fS1; + T pSH15 = -0.5900435899266435f * fC2; + T pSH9 = -0.5900435899266435f * fS2; + + result += + pSH9 * coeffs[9 * 3 + c] + pSH10 * coeffs[10 * 3 + c] + + pSH11 * coeffs[11 * 3 + c] + pSH12 * coeffs[12 * 3 + c] + + pSH13 * coeffs[13 * 3 + c] + pSH14 * coeffs[14 * 3 + c] + + pSH15 * coeffs[15 * 3 + c]; + + if (degree >= 4) { + T fTmp0D = + z * (-4.683325804901025f * z2 + 2.007139630671868f); + T fTmp1C = 3.31161143515146f * z2 - 0.47308734787878f; + T fTmp2B = -1.770130769779931f * z; + T fC3 = x * fC2 - y * fS2; + T fS3 = x * fS2 + y * fC2; + T pSH20 = + (1.984313483298443f * z * pSH12 - + 1.006230589874905f * pSH6); + T pSH21 = fTmp0D * x; + T pSH19 = fTmp0D * y; + T pSH22 = fTmp1C * fC1; + T pSH18 = fTmp1C * fS1; + T pSH23 = fTmp2B * fC2; + T pSH17 = fTmp2B * fS2; + T pSH24 = 0.6258357354491763f * fC3; + T pSH16 = 0.6258357354491763f * fS3; + + result += pSH16 * coeffs[16 * 3 + c] + + pSH17 * coeffs[17 * 3 + c] + + pSH18 * coeffs[18 * 3 + c] + + pSH19 * coeffs[19 * 3 + c] + + pSH20 * coeffs[20 * 3 + c] + + pSH21 * coeffs[21 * 3 + c] + + pSH22 * coeffs[22 * 3 + c] + + pSH23 * coeffs[23 * 3 + c] + + pSH24 * coeffs[24 * 3 + c]; + } + } + } + } + + colors[c] = result; +} + +template +inline void sh_coeffs_to_color_fast_vjp( + const uint32_t degree, // degree of SH to be evaluated + const uint32_t c, // color channel + const vec3 &dir, // [3] + const T *coeffs, // [K, 3] + const T *v_colors, // [3] + // output + T *v_coeffs, // [K, 3] + vec3 *v_dir // [3] optional +) { + T v_colors_local = v_colors[c]; + + v_coeffs[c] = 0.2820947917738781f * v_colors_local; + if (degree < 1) { + return; + } + T inorm = sycl::rsqrt(dir.x * dir.x + dir.y * dir.y + dir.z * dir.z); + T x = dir.x * inorm; + T y = dir.y * inorm; + T z = dir.z * inorm; + T v_x = 0.f, v_y = 0.f, v_z = 0.f; + + v_coeffs[1 * 3 + c] = -0.48860251190292f * y * v_colors_local; + v_coeffs[2 * 3 + c] = 0.48860251190292f * z * v_colors_local; + v_coeffs[3 * 3 + c] = -0.48860251190292f * x * v_colors_local; + + if (v_dir != nullptr) { + v_x += -0.48860251190292f * coeffs[3 * 3 + c] * v_colors_local; + v_y += -0.48860251190292f * coeffs[1 * 3 + c] * v_colors_local; + v_z += 0.48860251190292f * coeffs[2 * 3 + c] * v_colors_local; + } + if (degree < 2) { + if (v_dir != nullptr) { + vec3 dir_n = vec3(x, y, z); + vec3 v_dir_n = vec3(v_x, v_y, v_z); + vec3 v_d = (v_dir_n - glm::dot(v_dir_n, dir_n) * dir_n) * inorm; + + v_dir->x = v_d.x; + v_dir->y = v_d.y; + v_dir->z = v_d.z; + } + return; + } + + T z2 = z * z; + T fTmp0B = -1.092548430592079f * z; + T fC1 = x * x - y * y; + T fS1 = 2.f * x * y; + T pSH6 = (0.9461746957575601f * z2 - 0.3153915652525201f); + T pSH7 = fTmp0B * x; + T pSH5 = fTmp0B * y; + T pSH8 = 0.5462742152960395f * fC1; + T pSH4 = 0.5462742152960395f * fS1; + v_coeffs[4 * 3 + c] = pSH4 * v_colors_local; + v_coeffs[5 * 3 + c] = pSH5 * v_colors_local; + v_coeffs[6 * 3 + c] = pSH6 * v_colors_local; + v_coeffs[7 * 3 + c] = pSH7 * v_colors_local; + v_coeffs[8 * 3 + c] = pSH8 * v_colors_local; + + T fTmp0B_z, fC1_x, fC1_y, fS1_x, fS1_y, pSH6_z, pSH7_x, pSH7_z, pSH5_y, + pSH5_z, pSH8_x, pSH8_y, pSH4_x, pSH4_y; + if (v_dir != nullptr) { + fTmp0B_z = -1.092548430592079f; + fC1_x = 2.f * x; + fC1_y = -2.f * y; + fS1_x = 2.f * y; + fS1_y = 2.f * x; + pSH6_z = 2.f * 0.9461746957575601f * z; + pSH7_x = fTmp0B; + pSH7_z = fTmp0B_z * x; + pSH5_y = fTmp0B; + pSH5_z = fTmp0B_z * y; + pSH8_x = 0.5462742152960395f * fC1_x; + pSH8_y = 0.5462742152960395f * fC1_y; + pSH4_x = 0.5462742152960395f * fS1_x; + pSH4_y = 0.5462742152960395f * fS1_y; + + v_x += v_colors_local * + (pSH4_x * coeffs[4 * 3 + c] + pSH8_x * coeffs[8 * 3 + c] + + pSH7_x * coeffs[7 * 3 + c]); + v_y += v_colors_local * + (pSH4_y * coeffs[4 * 3 + c] + pSH8_y * coeffs[8 * 3 + c] + + pSH5_y * coeffs[5 * 3 + c]); + v_z += v_colors_local * + (pSH6_z * coeffs[6 * 3 + c] + pSH7_z * coeffs[7 * 3 + c] + + pSH5_z * coeffs[5 * 3 + c]); + } + + if (degree < 3) { + if (v_dir != nullptr) { + vec3 dir_n = vec3(x, y, z); + vec3 v_dir_n = vec3(v_x, v_y, v_z); + vec3 v_d = (v_dir_n - glm::dot(v_dir_n, dir_n) * dir_n) * inorm; + + v_dir->x = v_d.x; + v_dir->y = v_d.y; + v_dir->z = v_d.z; + } + return; + } + + T fTmp0C = -2.285228997322329f * z2 + 0.4570457994644658f; + T fTmp1B = 1.445305721320277f * z; + T fC2 = x * fC1 - y * fS1; + T fS2 = x * fS1 + y * fC1; + T pSH12 = z * (1.865881662950577f * z2 - 1.119528997770346f); + T pSH13 = fTmp0C * x; + T pSH11 = fTmp0C * y; + T pSH14 = fTmp1B * fC1; + T pSH10 = fTmp1B * fS1; + T pSH15 = -0.5900435899266435f * fC2; + T pSH9 = -0.5900435899266435f * fS2; + v_coeffs[9 * 3 + c] = pSH9 * v_colors_local; + v_coeffs[10 * 3 + c] = pSH10 * v_colors_local; + v_coeffs[11 * 3 + c] = pSH11 * v_colors_local; + v_coeffs[12 * 3 + c] = pSH12 * v_colors_local; + v_coeffs[13 * 3 + c] = pSH13 * v_colors_local; + v_coeffs[14 * 3 + c] = pSH14 * v_colors_local; + v_coeffs[15 * 3 + c] = pSH15 * v_colors_local; + + T fTmp0C_z, fTmp1B_z, fC2_x, fC2_y, fS2_x, fS2_y, pSH12_z, pSH13_x, pSH13_z, + pSH11_y, pSH11_z, pSH14_x, pSH14_y, pSH14_z, pSH10_x, pSH10_y, pSH10_z, + pSH15_x, pSH15_y, pSH9_x, pSH9_y; + if (v_dir != nullptr) { + fTmp0C_z = -2.285228997322329f * 2.f * z; + fTmp1B_z = 1.445305721320277f; + fC2_x = fC1 + x * fC1_x - y * fS1_x; + fC2_y = x * fC1_y - fS1 - y * fS1_y; + fS2_x = fS1 + x * fS1_x + y * fC1_x; + fS2_y = x * fS1_y + fC1 + y * fC1_y; + pSH12_z = 3.f * 1.865881662950577f * z2 - 1.119528997770346f; + pSH13_x = fTmp0C; + pSH13_z = fTmp0C_z * x; + pSH11_y = fTmp0C; + pSH11_z = fTmp0C_z * y; + pSH14_x = fTmp1B * fC1_x; + pSH14_y = fTmp1B * fC1_y; + pSH14_z = fTmp1B_z * fC1; + pSH10_x = fTmp1B * fS1_x; + pSH10_y = fTmp1B * fS1_y; + pSH10_z = fTmp1B_z * fS1; + pSH15_x = -0.5900435899266435f * fC2_x; + pSH15_y = -0.5900435899266435f * fC2_y; + pSH9_x = -0.5900435899266435f * fS2_x; + pSH9_y = -0.5900435899266435f * fS2_y; + + v_x += v_colors_local * + (pSH9_x * coeffs[9 * 3 + c] + pSH15_x * coeffs[15 * 3 + c] + + pSH10_x * coeffs[10 * 3 + c] + pSH14_x * coeffs[14 * 3 + c] + + pSH13_x * coeffs[13 * 3 + c]); + + v_y += v_colors_local * + (pSH9_y * coeffs[9 * 3 + c] + pSH15_y * coeffs[15 * 3 + c] + + pSH10_y * coeffs[10 * 3 + c] + pSH14_y * coeffs[14 * 3 + c] + + pSH11_y * coeffs[11 * 3 + c]); + + v_z += v_colors_local * + (pSH12_z * coeffs[12 * 3 + c] + pSH13_z * coeffs[13 * 3 + c] + + pSH11_z * coeffs[11 * 3 + c] + pSH14_z * coeffs[14 * 3 + c] + + pSH10_z * coeffs[10 * 3 + c]); + } + + if (degree < 4) { + if (v_dir != nullptr) { + vec3 dir_n = vec3(x, y, z); + vec3 v_dir_n = vec3(v_x, v_y, v_z); + vec3 v_d = (v_dir_n - glm::dot(v_dir_n, dir_n) * dir_n) * inorm; + + v_dir->x = v_d.x; + v_dir->y = v_d.y; + v_dir->z = v_d.z; + } + return; + } + + T fTmp0D = z * (-4.683325804901025f * z2 + 2.007139630671868f); + T fTmp1C = 3.31161143515146f * z2 - 0.47308734787878f; + T fTmp2B = -1.770130769779931f * z; + T fC3 = x * fC2 - y * fS2; + T fS3 = x * fS2 + y * fC2; + T pSH20 = (1.984313483298443f * z * pSH12 + -1.006230589874905f * pSH6); + T pSH21 = fTmp0D * x; + T pSH19 = fTmp0D * y; + T pSH22 = fTmp1C * fC1; + T pSH18 = fTmp1C * fS1; + T pSH23 = fTmp2B * fC2; + T pSH17 = fTmp2B * fS2; + T pSH24 = 0.6258357354491763f * fC3; + T pSH16 = 0.6258357354491763f * fS3; + v_coeffs[16 * 3 + c] = pSH16 * v_colors_local; + v_coeffs[17 * 3 + c] = pSH17 * v_colors_local; + v_coeffs[18 * 3 + c] = pSH18 * v_colors_local; + v_coeffs[19 * 3 + c] = pSH19 * v_colors_local; + v_coeffs[20 * 3 + c] = pSH20 * v_colors_local; + v_coeffs[21 * 3 + c] = pSH21 * v_colors_local; + v_coeffs[22 * 3 + c] = pSH22 * v_colors_local; + v_coeffs[23 * 3 + c] = pSH23 * v_colors_local; + v_coeffs[24 * 3 + c] = pSH24 * v_colors_local; + + T fTmp0D_z, fTmp1C_z, fTmp2B_z, fC3_x, fC3_y, fS3_x, fS3_y, pSH20_z, + pSH21_x, pSH21_z, pSH19_y, pSH19_z, pSH22_x, pSH22_y, pSH22_z, pSH18_x, + pSH18_y, pSH18_z, pSH23_x, pSH23_y, pSH23_z, pSH17_x, pSH17_y, pSH17_z, + pSH24_x, pSH24_y, pSH16_x, pSH16_y; + if (v_dir != nullptr) { + fTmp0D_z = 3.f * -4.683325804901025f * z2 + 2.007139630671868f; + fTmp1C_z = 2.f * 3.31161143515146f * z; + fTmp2B_z = -1.770130769779931f; + fC3_x = fC2 + x * fC2_x - y * fS2_x; + fC3_y = x * fC2_y - fS2 - y * fS2_y; + fS3_x = fS2 + y * fC2_x + x * fS2_x; + fS3_y = x * fS2_y + fC2 + y * fC2_y; + pSH20_z = 1.984313483298443f * (pSH12 + z * pSH12_z) + + -1.006230589874905f * pSH6_z; + pSH21_x = fTmp0D; + pSH21_z = fTmp0D_z * x; + pSH19_y = fTmp0D; + pSH19_z = fTmp0D_z * y; + pSH22_x = fTmp1C * fC1_x; + pSH22_y = fTmp1C * fC1_y; + pSH22_z = fTmp1C_z * fC1; + pSH18_x = fTmp1C * fS1_x; + pSH18_y = fTmp1C * fS1_y; + pSH18_z = fTmp1C_z * fS1; + pSH23_x = fTmp2B * fC2_x; + pSH23_y = fTmp2B * fC2_y; + pSH23_z = fTmp2B_z * fC2; + pSH17_x = fTmp2B * fS2_x; + pSH17_y = fTmp2B * fS2_y; + pSH17_z = fTmp2B_z * fS2; + pSH24_x = 0.6258357354491763f * fC3_x; + pSH24_y = 0.6258357354491763f * fC3_y; + pSH16_x = 0.6258357354491763f * fS3_x; + pSH16_y = 0.6258357354491763f * fS3_y; + + v_x += v_colors_local * + (pSH16_x * coeffs[16 * 3 + c] + pSH24_x * coeffs[24 * 3 + c] + + pSH17_x * coeffs[17 * 3 + c] + pSH23_x * coeffs[23 * 3 + c] + + pSH18_x * coeffs[18 * 3 + c] + pSH22_x * coeffs[22 * 3 + c] + + pSH21_x * coeffs[21 * 3 + c]); + v_y += v_colors_local * + (pSH16_y * coeffs[16 * 3 + c] + pSH24_y * coeffs[24 * 3 + c] + + pSH17_y * coeffs[17 * 3 + c] + pSH23_y * coeffs[23 * 3 + c] + + pSH18_y * coeffs[18 * 3 + c] + pSH22_y * coeffs[22 * 3 + c] + + pSH19_y * coeffs[19 * 3 + c]); + v_z += v_colors_local * + (pSH20_z * coeffs[20 * 3 + c] + pSH21_z * coeffs[21 * 3 + c] + + pSH19_z * coeffs[19 * 3 + c] + pSH22_z * coeffs[22 * 3 + c] + + pSH18_z * coeffs[18 * 3 + c] + pSH23_z * coeffs[23 * 3 + c] + + pSH17_z * coeffs[17 * 3 + c]); + + vec3 dir_n = vec3(x, y, z); + vec3 v_dir_n = vec3(v_x, v_y, v_z); + vec3 v_d = (v_dir_n - glm::dot(v_dir_n, dir_n) * dir_n) * inorm; + + v_dir->x = v_d.x; + v_dir->y = v_d.y; + v_dir->z = v_d.z; + } +} \ No newline at end of file diff --git a/gsplat/sycl/include/transform.hpp b/gsplat/sycl/include/transform.hpp new file mode 100644 index 00000000..917c8a54 --- /dev/null +++ b/gsplat/sycl/include/transform.hpp @@ -0,0 +1,66 @@ +#pragma once + +#include "types.hpp" + +template +inline void pos_world_to_cam( + // [R, t] is the world-to-camera transformation + const mat3 R, + const vec3 t, + const vec3 p, + vec3 &p_c +) { + p_c = R * p + t; +} + +template +inline void pos_world_to_cam_vjp( + // fwd inputs + const mat3 R, + const vec3 t, + const vec3 p, + // grad outputs + const vec3 v_p_c, + // grad inputs + mat3 &v_R, + vec3 &v_t, + vec3 &v_p +) { + // for D = W * X, G = df/dD + // df/dW = G * XT, df/dX = WT * G + v_R += glm::outerProduct(v_p_c, p); + v_t += v_p_c; + v_p += glm::transpose(R) * v_p_c; +} + +template +inline void covar_world_to_cam( + // [R, t] is the world-to-camera transformation + const mat3 R, + const mat3 covar, + mat3 &covar_c +) { + covar_c = R * covar * glm::transpose(R); +} + +template +inline void covar_world_to_cam_vjp( + // fwd inputs + const mat3 R, + const mat3 covar, + // grad outputs + const mat3 v_covar_c, + // grad inputs + mat3 &v_R, + mat3 &v_covar +) { + // for D = W * X * WT, G = df/dD + // df/dX = WT * G * W + // df/dW + // = G * (X * WT)T + ((W * X)T * G)T + // = G * W * XT + (XT * WT * G)T + // = G * W * XT + GT * W * X + v_R += v_covar_c * R * glm::transpose(covar) + + glm::transpose(v_covar_c) * R * covar; + v_covar += glm::transpose(R) * v_covar_c * R; +} \ No newline at end of file diff --git a/gsplat/sycl/include/types.hpp b/gsplat/sycl/include/types.hpp new file mode 100644 index 00000000..109b5536 --- /dev/null +++ b/gsplat/sycl/include/types.hpp @@ -0,0 +1,17 @@ +#pragma once + +#include + +template using vec2 = glm::vec<2, T>; + +template using vec3 = glm::vec<3, T>; + +template using vec4 = glm::vec<4, T>; + +template using mat2 = glm::mat<2, 2, T>; + +template using mat3 = glm::mat<3, 3, T>; + +template using mat4 = glm::mat<4, 4, T>; + +template using mat3x2 = glm::mat<3, 2, T>; \ No newline at end of file diff --git a/gsplat/sycl/include/utils.hpp b/gsplat/sycl/include/utils.hpp new file mode 100644 index 00000000..0bdd1059 --- /dev/null +++ b/gsplat/sycl/include/utils.hpp @@ -0,0 +1,72 @@ +#pragma once + +#include "types.hpp" + +#include + +template inline T inverse(const mat2 M, mat2 &Minv) { + T det = M[0][0] * M[1][1] - M[0][1] * M[1][0]; + if (det <= 0.f) { + return det; + } + T invDet = 1.f / det; + Minv[0][0] = M[1][1] * invDet; + Minv[0][1] = -M[0][1] * invDet; + Minv[1][0] = Minv[0][1]; + Minv[1][1] = M[0][0] * invDet; + return det; +} + +template +inline void inverse_vjp(const T Minv, const T v_Minv, T &v_M) { + // P = M^-1 + // df/dM = -P * df/dP * P + v_M += -Minv * v_Minv * Minv; +} + +template +inline T add_blur(const T eps2d, mat2 &covar, T &compensation) { + T det_orig = covar[0][0] * covar[1][1] - covar[0][1] * covar[1][0]; + covar[0][0] += eps2d; + covar[1][1] += eps2d; + T det_blur = covar[0][0] * covar[1][1] - covar[0][1] * covar[1][0]; + compensation = + sycl::sqrt(sycl::max(static_cast(0), det_orig / det_blur)); + return det_blur; +} + +template +inline void add_blur_vjp( + const T eps2d, + const mat2 conic_blur, + const T compensation, + const T v_compensation, + mat2 &v_covar +) { + // comp = sqrt(det(covar) / det(covar_blur)) + + // d [det(M)] / d M = adj(M) + // d [det(M + aI)] / d M = adj(M + aI) = adj(M) + a * I + // d [det(M) / det(M + aI)] / d M + // = (det(M + aI) * adj(M) - det(M) * adj(M + aI)) / (det(M + aI))^2 + // = adj(M) / det(M + aI) - adj(M + aI) / det(M + aI) * comp^2 + // = (adj(M) - adj(M + aI) * comp^2) / det(M + aI) + // given that adj(M + aI) = adj(M) + a * I + // = (adj(M + aI) - aI - adj(M + aI) * comp^2) / det(M + aI) + // given that adj(M) / det(M) = inv(M) + // = (1 - comp^2) * inv(M + aI) - aI / det(M + aI) + // given det(inv(M)) = 1 / det(M) + // = (1 - comp^2) * inv(M + aI) - aI * det(inv(M + aI)) + // = (1 - comp^2) * conic_blur - aI * det(conic_blur) + + T det_conic_blur = conic_blur[0][0] * conic_blur[1][1] - + conic_blur[0][1] * conic_blur[1][0]; + T v_sqr_comp = v_compensation * 0.5f / (compensation + 1e-6f); + T one_minus_sqr_comp = 1 - compensation * compensation; + v_covar[0][0] += v_sqr_comp * (one_minus_sqr_comp * conic_blur[0][0] - + eps2d * det_conic_blur); + v_covar[0][1] += v_sqr_comp * (one_minus_sqr_comp * conic_blur[0][1]); + v_covar[1][0] += v_sqr_comp * (one_minus_sqr_comp * conic_blur[1][0]); + v_covar[1][1] += v_sqr_comp * (one_minus_sqr_comp * conic_blur[1][1] - + eps2d * det_conic_blur); +} \ No newline at end of file diff --git a/gsplat/sycl/src/adam.cpp b/gsplat/sycl/src/adam.cpp new file mode 100644 index 00000000..7abfd6d8 --- /dev/null +++ b/gsplat/sycl/src/adam.cpp @@ -0,0 +1,23 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +void adam( + at::Tensor ¶m, // [..., D] + const at::Tensor ¶m_grad, // [..., D] + at::Tensor &exp_avg, // [..., D] + at::Tensor &exp_avg_sq, // [..., D] + const at::optional valid, // [...] + const float lr, + const float b1, + const float b2, + const float eps +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/intersect_offset.cpp b/gsplat/sycl/src/intersect_offset.cpp new file mode 100644 index 00000000..0f6c7da6 --- /dev/null +++ b/gsplat/sycl/src/intersect_offset.cpp @@ -0,0 +1,57 @@ +#include + +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/IsectOffsetEncodeKernel.hpp" + +namespace gsplat::xpu { + +at::Tensor intersect_offset( + const at::Tensor isect_ids, // [n_isects] + const uint32_t I, + const uint32_t tile_width, + const uint32_t tile_height +) { + DEVICE_GUARD(isect_ids); + CHECK_INPUT(isect_ids); + const uint32_t C = I; + + auto options = isect_ids.options().dtype(at::kInt); + at::Tensor offsets = at::empty({C, tile_height, tile_width}, options); + + const uint32_t n_isects = isect_ids.size(0); + + if (n_isects > 0) { + const uint32_t n_tiles = tile_width * tile_height; + const uint32_t tile_n_bits = (uint32_t)floor(log2(n_tiles)) + 1; + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = + (n_isects + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + IsectOffsetEncodeKernel kernel( + n_isects, + isect_ids.data_ptr(), + C, + n_tiles, + tile_n_bits, + offsets.data_ptr() + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + } else { + offsets.fill_(0); + } + + return offsets; +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/intersect_tile.cpp b/gsplat/sycl/src/intersect_tile.cpp new file mode 100644 index 00000000..8d9e9f5c --- /dev/null +++ b/gsplat/sycl/src/intersect_tile.cpp @@ -0,0 +1,147 @@ +#include +#include + +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/IsectTilesKernel.hpp" + +namespace gsplat::xpu { + +std::tuple intersect_tile( + const at::Tensor means2d, // [..., C, N, 2] or [nnz, 2] + const at::Tensor radii, // [..., C, N] or [nnz] + const at::Tensor depths, // [..., C, N] or [nnz] + const at::optional image_ids, // [nnz] -> maps to camera_ids + const at::optional gaussian_ids, // [nnz] + const uint32_t I, // -> maps to C + const uint32_t tile_size, + const uint32_t tile_width, + const uint32_t tile_height, + const bool sort, + const bool segmented +) { + DEVICE_GUARD(means2d); + CHECK_INPUT(means2d); + CHECK_INPUT2(radii, means2d); + CHECK_INPUT2(depths, means2d); + if (image_ids.has_value()) + CHECK_INPUT2(image_ids.value(), means2d); + if (gaussian_ids.has_value()) + CHECK_INPUT2(gaussian_ids.value(), means2d); + + const bool packed = segmented; + const uint32_t C = I; + uint32_t N = 0; + uint32_t nnz = 0; + uint32_t total_elems = 0; + + if (packed) { + nnz = means2d.size(0); + total_elems = nnz; + TORCH_CHECK( + (image_ids.has_value()) && (gaussian_ids.has_value()), + "When segmented (packed) is set, image_ids and gaussian_ids " + "must be provided." + ); + } else { + N = means2d.size(-2); + total_elems = C * N; + } + + if (total_elems == 0) { + return std::make_tuple( + at::empty_like(depths, at::kInt), + at::empty({0}, at::kLong), + at::empty({0}, at::kInt) + ); + } + auto options = depths.options(); + at::Tensor tiles_per_gauss = + at::empty_like(depths, options.dtype(at::kInt)); + const uint32_t n_tiles = tile_width * tile_height; + const uint32_t tile_n_bits = (uint32_t)floor(log2(n_tiles)) + 1; + const uint32_t cam_n_bits = (uint32_t)floor(log2(C)) + 1; + TORCH_CHECK( + tile_n_bits + cam_n_bits <= 32, + "Not enough bits to encode camera and tile IDs." + ); + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + size_t numWorkGrps = + (total_elems + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e1 = d_queue.submit([&](sycl::handler &cgh) { + IsectTilesKernel kernel( + packed, + C, + N, + nnz, + packed ? image_ids.value().data_ptr() : nullptr, + packed ? gaussian_ids.value().data_ptr() : nullptr, + means2d.data_ptr(), + radii.data_ptr(), + depths.data_ptr(), + nullptr, // cum_tiles_per_gauss + tile_size, + tile_width, + tile_height, + tile_n_bits, + tiles_per_gauss.data_ptr(), + nullptr, // isect_ids + nullptr // flatten_ids + ); + cgh.parallel_for(range, kernel); + }); + e1.wait(); + + at::Tensor cum_tiles_per_gauss = + at::cumsum(tiles_per_gauss.view({-1}), 0, at::kLong); + int64_t n_isects = 0; + if (total_elems > 0) { + n_isects = cum_tiles_per_gauss.slice(0, -1).item(); + } + + at::Tensor isect_ids = at::empty({n_isects}, options.dtype(at::kLong)); + at::Tensor flatten_ids = at::empty({n_isects}, options.dtype(at::kInt)); + + if (n_isects > 0) { + auto e2 = d_queue.submit([&](sycl::handler &cgh) { + IsectTilesKernel kernel( + packed, + C, + N, + nnz, + packed ? image_ids.value().data_ptr() : nullptr, + packed ? gaussian_ids.value().data_ptr() : nullptr, + means2d.data_ptr(), + radii.data_ptr(), + depths.data_ptr(), + cum_tiles_per_gauss.data_ptr(), + tile_size, + tile_width, + tile_height, + tile_n_bits, + nullptr, // tiles_per_gauss + isect_ids.data_ptr(), + flatten_ids.data_ptr() + ); + cgh.parallel_for(range, kernel); + }); + e2.wait(); + } + + if (n_isects > 0 && sort) { + auto [sorted_isect_ids, sort_indices] = at::sort(isect_ids); + isect_ids = sorted_isect_ids; + flatten_ids = flatten_ids.index_select(0, sort_indices); + } + + return std::make_tuple(tiles_per_gauss, isect_ids, flatten_ids); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/null.cpp b/gsplat/sycl/src/null.cpp new file mode 100644 index 00000000..bfcecb03 --- /dev/null +++ b/gsplat/sycl/src/null.cpp @@ -0,0 +1,13 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +at::Tensor null(const at::Tensor input) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp new file mode 100644 index 00000000..796b7834 --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp @@ -0,0 +1,121 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/Projection2DGSFusedBwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple +projection_2dgs_fused_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + // fwd outputs + const at::Tensor radii, // [..., C, N, 2] + const at::Tensor ray_transforms, // [..., C, N, 3, 3] + // grad outputs + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_depths, // [..., C, N] + const at::Tensor v_normals, // [..., C, N, 3] + const at::Tensor v_ray_transforms, // [..., C, N, 3, 3] + const bool viewmats_requires_grad +) { + DEVICE_GUARD(means); + CHECK_INPUT(means); + CHECK_INPUT2(quats, means); + CHECK_INPUT2(scales, means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); + CHECK_INPUT2(radii, means); + CHECK_INPUT2(ray_transforms, means); + CHECK_INPUT2(v_means2d, means); + CHECK_INPUT2(v_depths, means); + CHECK_INPUT2(v_normals, means); + CHECK_INPUT2(v_ray_transforms, means); + + TORCH_CHECK( + means.dim() >= 2, "means must have at least 2 dimensions [..., N, 3]" + ); + TORCH_CHECK( + quats.dim() >= 2, "quats must have at least 2 dimensions [..., N, 4]" + ); + TORCH_CHECK( + scales.dim() >= 2, "scales must have at least 2 dimensions [..., N, 3]" + ); + TORCH_CHECK( + viewmats.dim() >= 3, + "viewmats must have at least 3 dimensions [..., C, 4, 4]" + ); + + const uint32_t N = means.size(-2); // number of gaussians + const uint32_t C = viewmats.size(-3); // number of cameras + const uint32_t B = means.numel() / (N * 3); // number of batches + const int64_t n_elements = B * C * N; + + auto options = means.options(); + + // Initialize gradient tensors + at::Tensor v_means = at::zeros_like(means); + at::Tensor v_quats = at::zeros_like(quats); + at::Tensor v_scales = at::zeros_like(scales); + at::Tensor v_viewmats = + viewmats_requires_grad ? at::zeros_like(viewmats) : at::Tensor(); + + if (n_elements == 0) { + // Skip kernel launch if there are no elements + return std::make_tuple(v_means, v_quats, v_scales, v_viewmats); + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + auto num_work_groups = + (n_elements + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> local_range(GSPLAT_N_THREADS); + sycl::range<1> global_range(num_work_groups * GSPLAT_N_THREADS); + + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_2dgs_fused_bwd", + [&] { + auto e = d_queue.submit([&](sycl::handler &cgh) { + Projection2DGSFusedBwdKernel kernel( + B, + C, + N, + means.data_ptr(), + quats.data_ptr(), + scales.data_ptr(), + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + radii.data_ptr(), + ray_transforms.data_ptr(), + v_means2d.data_ptr(), + v_depths.data_ptr(), + v_normals.data_ptr(), + v_ray_transforms.data_ptr(), + v_means.data_ptr(), + v_quats.data_ptr(), + v_scales.data_ptr(), + viewmats_requires_grad ? v_viewmats.data_ptr() + : nullptr + ); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel + ); + }); + e.wait(); + } + ); + + return std::make_tuple(v_means, v_quats, v_scales, v_viewmats); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp new file mode 100644 index 00000000..c97b8423 --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp @@ -0,0 +1,130 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/Projection2DGSFusedFwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_2dgs_fused_fwd( + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip +) { + DEVICE_GUARD(means); + CHECK_INPUT(means); + CHECK_INPUT2(quats, means); + CHECK_INPUT2(scales, means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); + + TORCH_CHECK( + means.dim() >= 2, "means must have at least 2 dimensions [..., N, 3]" + ); + TORCH_CHECK( + quats.dim() >= 2, "quats must have at least 2 dimensions [..., N, 4]" + ); + TORCH_CHECK( + scales.dim() >= 2, "scales must have at least 2 dimensions [..., N, 3]" + ); + TORCH_CHECK( + viewmats.dim() >= 3, + "viewmats must have at least 3 dimensions [..., C, 4, 4]" + ); + TORCH_CHECK( + Ks.dim() >= 3, "Ks must have at least 3 dimensions [..., C, 3, 3]" + ); + + const uint32_t N = means.size(-2); // number of gaussians + const uint32_t C = viewmats.size(-3); // number of cameras + const uint32_t B = means.numel() / (N * 3); // number of batches + const int64_t n_elements = B * C * N; + + auto options = means.options(); + at::DimVector batch_dims(means.sizes().slice(0, means.dim() - 2)); + + // Output shape: [..., C, N] + at::DimVector out_shape_cn = batch_dims; + out_shape_cn.insert(out_shape_cn.end(), {C, N}); + + // Output shape: [..., C, N, 2] + at::DimVector out_shape_cn2 = batch_dims; + out_shape_cn2.insert(out_shape_cn2.end(), {C, N, 2}); + + // Output shape: [..., C, N, 3] + at::DimVector out_shape_cn3 = batch_dims; + out_shape_cn3.insert(out_shape_cn3.end(), {C, N, 3}); + + // Output shape: [..., C, N, 3, 3] + at::DimVector out_shape_cn33 = batch_dims; + out_shape_cn33.insert(out_shape_cn33.end(), {C, N, 3, 3}); + + at::Tensor radii = at::empty(out_shape_cn2, options.dtype(at::kInt)); + at::Tensor means2d = at::empty(out_shape_cn2, options); + at::Tensor depths = at::empty(out_shape_cn, options); + at::Tensor ray_transforms = at::empty(out_shape_cn33, options); + at::Tensor normals = at::empty(out_shape_cn3, options); + + if (n_elements == 0) { + // Skip kernel launch if there are no elements + return std::make_tuple(radii, means2d, depths, ray_transforms, normals); + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + auto num_work_groups = + (n_elements + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> local_range(GSPLAT_N_THREADS); + sycl::range<1> global_range(num_work_groups * GSPLAT_N_THREADS); + + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_2dgs_fused_fwd", + [&] { + auto e = d_queue.submit([&](sycl::handler &cgh) { + Projection2DGSFusedFwdKernel kernel( + B, + C, + N, + means.data_ptr(), + quats.data_ptr(), + scales.data_ptr(), + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + near_plane, + far_plane, + radius_clip, + radii.data_ptr(), + means2d.data_ptr(), + depths.data_ptr(), + ray_transforms.data_ptr(), + normals.data_ptr() + ); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel + ); + }); + e.wait(); + } + ); + + return std::make_tuple(radii, means2d, depths, ray_transforms, normals); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp b/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp new file mode 100644 index 00000000..bf36b64f --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp @@ -0,0 +1,35 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple +projection_2dgs_packed_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + // fwd outputs + const at::Tensor batch_ids, // [nnz] + const at::Tensor camera_ids, // [nnz] + const at::Tensor gaussian_ids, // [nnz] + const at::Tensor ray_transforms, // [nnz, 3, 3] + // grad outputs + const at::Tensor v_means2d, // [nnz, 2] + const at::Tensor v_depths, // [nnz] + const at::Tensor v_ray_transforms, // [nnz, 3, 3] + const at::Tensor v_normals, // [nnz, 3] + const bool viewmats_requires_grad, + const bool sparse_grad +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp new file mode 100644 index 00000000..45c0005d --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp @@ -0,0 +1,34 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_2dgs_packed_fwd( + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float near_plane, + const float far_plane, + const float radius_clip +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp new file mode 100644 index 00000000..6510bad2 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp @@ -0,0 +1,134 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/FullyFusedProjectionBwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple +projection_ewa_3dgs_fused_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] optional + const at::optional quats, // [..., N, 4] optional + const at::optional scales, // [..., N, 3] optional + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const CameraModelType camera_model, + // fwd outputs + const at::Tensor radii, // [..., C, N, 2] + const at::Tensor conics, // [..., C, N, 3] + const at::optional compensations, // [..., C, N] optional + // grad outputs + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_depths, // [..., C, N] + const at::Tensor v_conics, // [..., C, N, 3] + const at::optional v_compensations, // [..., C, N] optional + const bool viewmats_requires_grad +) { + DEVICE_GUARD(means); + // Input validation + CHECK_INPUT(means); + if (covars.has_value()) + CHECK_INPUT2(covars.value(), means); + if (quats.has_value()) + CHECK_INPUT2(quats.value(), means); + if (scales.has_value()) + CHECK_INPUT2(scales.value(), means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); + CHECK_INPUT2(radii, means); + CHECK_INPUT2(conics, means); + if (compensations.has_value()) + CHECK_INPUT2(compensations.value(), means); + CHECK_INPUT2(v_means2d, means); + CHECK_INPUT2(v_depths, means); + CHECK_INPUT2(v_conics, means); + if (v_compensations.has_value()) + CHECK_INPUT2(v_compensations.value(), means); + + // Dimensions + const uint32_t N = means.size(-2); + const uint32_t C = viewmats.size(-3); + const uint32_t B = means.numel() / (N * 3); + const int64_t n_elements = B * C * N; + + // Create gradient tensors, initialized to zero + at::Tensor v_means = at::zeros_like(means); + at::Tensor v_covars = covars.has_value() ? at::zeros_like(covars.value()) + : at::empty({0}, means.options()); + at::Tensor v_quats = quats.has_value() ? at::zeros_like(quats.value()) + : at::empty({0}, means.options()); + at::Tensor v_scales = scales.has_value() ? at::zeros_like(scales.value()) + : at::empty({0}, means.options()); + at::Tensor v_viewmats = viewmats_requires_grad + ? at::zeros_like(viewmats) + : at::empty({0}, means.options()); + + if (n_elements > 0) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + auto num_work_groups = + (n_elements + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> local_range(GSPLAT_N_THREADS); + sycl::range<1> global_range(num_work_groups * GSPLAT_N_THREADS); + + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_ewa_3dgs_fused_bwd", + [&] { + auto e = d_queue.submit([&](sycl::handler &cgh) { + FullyFusedProjectionBwdKernel kernel( + B, + C, + N, + means.data_ptr(), + covars.has_value() ? covars.value().data_ptr() + : nullptr, + quats.has_value() ? quats.value().data_ptr() + : nullptr, + scales.has_value() ? scales.value().data_ptr() + : nullptr, + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + eps2d, + camera_model, + radii.data_ptr(), + conics.data_ptr(), + compensations.has_value() + ? compensations.value().data_ptr() + : nullptr, + v_means2d.data_ptr(), + v_depths.data_ptr(), + v_conics.data_ptr(), + v_compensations.has_value() + ? v_compensations.value().data_ptr() + : nullptr, + v_means.data_ptr(), + covars.has_value() ? v_covars.data_ptr() + : nullptr, + quats.has_value() ? v_quats.data_ptr() + : nullptr, + scales.has_value() ? v_scales.data_ptr() + : nullptr, + viewmats_requires_grad ? v_viewmats.data_ptr() + : nullptr + ); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel + ); + }); + e.wait(); + } + ); + } + + return std::make_tuple(v_means, v_covars, v_quats, v_scales, v_viewmats); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp new file mode 100644 index 00000000..9dbeeea9 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp @@ -0,0 +1,134 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/FullyFusedProjectionFwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_ewa_3dgs_fused_fwd( + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] optional + const at::optional quats, // [..., N, 4] optional + const at::optional scales, // [..., N, 3] optional + const at::optional opacities, // [..., N] optional + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip, + const bool calc_compensations, + const CameraModelType camera_model +) { + DEVICE_GUARD(means); + // Input validation + CHECK_INPUT(means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); + if (covars.has_value()) + CHECK_INPUT2(covars.value(), means); + if (quats.has_value()) + CHECK_INPUT2(quats.value(), means); + if (scales.has_value()) + CHECK_INPUT2(scales.value(), means); + if (opacities.has_value()) + CHECK_INPUT2(opacities.value(), means); + + TORCH_CHECK( + means.dim() >= 2, "means must have at least 2 dimensions [..., N, 3]" + ); + TORCH_CHECK( + viewmats.dim() >= 3, + "viewmats must have at least 3 dimensions [..., C, 4, 4]" + ); + + const uint32_t N = means.size(-2); // number of gaussians + const uint32_t C = viewmats.size(-3); // number of cameras + const uint32_t B = means.numel() / (N * 3); // number of batches + const int64_t n_elements = B * C * N; + + auto options = means.options(); + at::DimVector batch_dims(means.sizes().slice(0, means.dim() - 2)); + + at::DimVector out_shape_cn = batch_dims; + out_shape_cn.insert(out_shape_cn.end(), {C, N}); + + at::DimVector out_shape_cn2 = batch_dims; + out_shape_cn2.insert(out_shape_cn2.end(), {C, N, 2}); + + at::DimVector out_shape_cn3 = batch_dims; + out_shape_cn3.insert(out_shape_cn3.end(), {C, N, 3}); + + at::Tensor radii = at::empty(out_shape_cn2, options.dtype(at::kInt)); + at::Tensor means2d = at::empty(out_shape_cn2, options); + at::Tensor depths = at::empty(out_shape_cn, options); + at::Tensor conics = at::empty(out_shape_cn3, options); + at::Tensor compensations = at::empty(out_shape_cn, options); + + if (n_elements > 0) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + const auto dev_id = + d_queue.get_device().get_info(); + + auto num_work_groups = + (n_elements + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> local_range(GSPLAT_N_THREADS); + sycl::range<1> global_range(num_work_groups * GSPLAT_N_THREADS); + + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_ewa_3dgs_fused_fwd", + [&] { + auto e = d_queue.submit([&](sycl::handler &cgh) { + FullyFusedProjectionFwdKernel kernel( + B, + C, + N, + means.data_ptr(), + covars.has_value() ? covars.value().data_ptr() + : nullptr, + quats.has_value() ? quats.value().data_ptr() + : nullptr, + scales.has_value() ? scales.value().data_ptr() + : nullptr, + opacities.has_value() + ? opacities.value().data_ptr() + : nullptr, + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + eps2d, + near_plane, + far_plane, + radius_clip, + camera_model, + radii.data_ptr(), + means2d.data_ptr(), + depths.data_ptr(), + conics.data_ptr(), + calc_compensations ? compensations.data_ptr() + : nullptr + ); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel + ); + }); + e.wait(); + } + ); + } + + return std::make_tuple(radii, means2d, depths, conics, compensations); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp new file mode 100644 index 00000000..45b84d2b --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp @@ -0,0 +1,153 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/PackedProjectionBwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple +projection_ewa_3dgs_packed_bwd( + // fwd inputs + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] + const at::optional quats, // [..., N, 4] + const at::optional scales, // [..., N, 3] + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const CameraModelType camera_model, + // fwd outputs + const at::Tensor batch_ids, // [nnz] + const at::Tensor camera_ids, // [nnz] + const at::Tensor gaussian_ids, // [nnz] + const at::Tensor conics, // [nnz, 3] + const at::optional compensations, // [nnz] optional + // grad outputs + const at::Tensor v_means2d, // [nnz, 2] + const at::Tensor v_depths, // [nnz] + const at::Tensor v_conics, // [nnz, 3] + const at::optional v_compensations, // [nnz] optional + const bool viewmats_requires_grad, + const bool sparse_grad +) { + DEVICE_GUARD(means); + // Input validation + CHECK_INPUT(means); + if (covars.has_value()) + CHECK_INPUT2(covars.value(), means); + if (quats.has_value()) + CHECK_INPUT2(quats.value(), means); + if (scales.has_value()) + CHECK_INPUT2(scales.value(), means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); + CHECK_INPUT2(batch_ids, means); + CHECK_INPUT2(camera_ids, means); + CHECK_INPUT2(gaussian_ids, means); + CHECK_INPUT2(conics, means); + if (compensations.has_value()) + CHECK_INPUT2(compensations.value(), means); + CHECK_INPUT2(v_means2d, means); + CHECK_INPUT2(v_depths, means); + CHECK_INPUT2(v_conics, means); + if (v_compensations.has_value()) + CHECK_INPUT2(v_compensations.value(), means); + + uint32_t N = means.size(-2); + uint32_t C = viewmats.size(-3); + uint32_t B = means.numel() / (N * 3); + uint32_t nnz = batch_ids.size(0); + + // Allocate output gradient tensors + at::Tensor v_means, v_covars, v_quats, v_scales, v_viewmats; + + if (sparse_grad) { + v_means = at::empty({(long)nnz, 3}, means.options()); + if (covars.has_value()) { + v_covars = at::empty({(long)nnz, 6}, covars.value().options()); + } else { + v_quats = at::empty({(long)nnz, 4}, quats.value().options()); + v_scales = at::empty({(long)nnz, 3}, scales.value().options()); + } + } else { + v_means = at::zeros_like(means); + if (covars.has_value()) { + v_covars = at::zeros_like(covars.value()); + } else { + v_quats = at::zeros_like(quats.value()); + v_scales = at::zeros_like(scales.value()); + } + } + + if (viewmats_requires_grad) { + v_viewmats = at::zeros_like(viewmats); + } + + if (nnz == 0) { + return std::make_tuple( + v_means, v_covars, v_quats, v_scales, v_viewmats + ); + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + sycl::range<1> local_range(256); + sycl::range<1> global_range( + (nnz + local_range[0] - 1) / local_range[0] * local_range[0] + ); + sycl::nd_range<1> range(global_range, local_range); + + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_ewa_3dgs_packed_bwd_kernel", + [&] { + PackedProjectionBwdKernel kernel( + B, + C, + N, + nnz, + means.data_ptr(), + covars.has_value() ? covars.value().data_ptr() + : nullptr, + covars.has_value() ? nullptr + : quats.value().data_ptr(), + covars.has_value() ? nullptr + : scales.value().data_ptr(), + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + (scalar_t)eps2d, + camera_model, + batch_ids.data_ptr(), + camera_ids.data_ptr(), + gaussian_ids.data_ptr(), + conics.data_ptr(), + compensations.has_value() + ? compensations.value().data_ptr() + : nullptr, + v_means2d.data_ptr(), + v_depths.data_ptr(), + v_conics.data_ptr(), + v_compensations.has_value() + ? v_compensations.value().data_ptr() + : nullptr, + sparse_grad, + v_means.data_ptr(), + covars.has_value() ? v_covars.data_ptr() : nullptr, + covars.has_value() ? nullptr : v_quats.data_ptr(), + covars.has_value() ? nullptr : v_scales.data_ptr(), + viewmats_requires_grad ? v_viewmats.data_ptr() + : nullptr + ); + auto e = d_queue.parallel_for(range, kernel); + e.wait(); + } + ); + + return std::make_tuple(v_means, v_covars, v_quats, v_scales, v_viewmats); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp new file mode 100644 index 00000000..168d8d46 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -0,0 +1,246 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/PackedProjectionFwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_ewa_3dgs_packed_fwd( + const at::Tensor means, // [..., N, 3] + const at::optional covars, // [..., N, 6] optional + const at::optional quats, // [..., N, 4] optional + const at::optional scales, // [..., N, 3] optional + const at::optional opacities, // [..., N] optional + const at::Tensor viewmats, // [..., C, 4, 4] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip, + const bool calc_compensations, + const CameraModelType camera_model +) { + DEVICE_GUARD(means); + // Input validation + CHECK_INPUT(means); + if (covars.has_value()) + CHECK_INPUT2(covars.value(), means); + if (quats.has_value()) + CHECK_INPUT2(quats.value(), means); + if (scales.has_value()) + CHECK_INPUT2(scales.value(), means); + if (opacities.has_value()) + CHECK_INPUT2(opacities.value(), means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); + + uint32_t N = means.size(-2); + uint32_t C = viewmats.size(-3); + uint32_t B = means.numel() / (N * 3); + + uint32_t nrows = B * C; + uint32_t ncols = N; + uint32_t blocks_per_row = (ncols + N_THREADS_PACKED - 1) / N_THREADS_PACKED; + uint32_t n_blocks = nrows * blocks_per_row; + + // Create empty outputs for the case where there's nothing to process + auto long_opts = means.options().dtype(at::kLong); + auto int_opts = means.options().dtype(at::kInt); + auto float_opts = means.options(); + + at::Tensor batch_ids = at::empty({0}, long_opts); + at::Tensor camera_ids = at::empty({0}, long_opts); + at::Tensor gaussian_ids = at::empty({0}, long_opts); + at::Tensor radii = at::empty({0, 2}, int_opts); + at::Tensor means2d = at::empty({0, 2}, float_opts); + at::Tensor depths = at::empty({0}, float_opts); + at::Tensor conics = at::empty({0, 3}, float_opts); + at::Tensor indptr = at::zeros({nrows + 1}, int_opts); + at::Tensor compensations = at::empty({0}, float_opts); + + if (B == 0 || C == 0 || N == 0) { + return std::make_tuple( + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + indptr, + compensations + ); + } + + // Allocate block_cnts as kInt, which the kernel expects. + at::Tensor block_cnts = at::empty({(long)n_blocks}, int_opts); + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + sycl::range<2> local_range(1, N_THREADS_PACKED); + sycl::range<2> global_range(nrows, blocks_per_row * N_THREADS_PACKED); + sycl::nd_range<2> range(global_range, local_range); + + // First pass: count visible Gaussians per block + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_ewa_3dgs_packed_fwd_kernel_pass1", + [&] { + d_queue + .parallel_for( + range, + PackedProjectionFwdKernel( + B, + C, + N, + means.data_ptr(), + covars.has_value() ? covars.value().data_ptr() + : nullptr, + quats.has_value() ? quats.value().data_ptr() + : nullptr, + scales.has_value() ? scales.value().data_ptr() + : nullptr, + opacities.has_value() + ? opacities.value().data_ptr() + : nullptr, + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + (scalar_t)eps2d, + (scalar_t)near_plane, + (scalar_t)far_plane, + (scalar_t)radius_clip, + camera_model, + nullptr, // block_accum + block_cnts.data_ptr(), + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr + ) + ) + .wait(); + } + ); + + // Perform inclusive scan on a kLong version of block_cnts to prevent + // overflow. + at::Tensor block_accum_inclusive = at::cumsum(block_cnts.to(at::kLong), 0); + + int64_t nnz = 0; + if (n_blocks > 0) { + nnz = block_accum_inclusive.index({-1}).item(); + } + + if (nnz == 0) { + return std::make_tuple( + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + indptr, + compensations + ); + } + + // Allocate final output tensors + batch_ids = at::empty({nnz}, long_opts); + camera_ids = at::empty({nnz}, long_opts); + gaussian_ids = at::empty({nnz}, long_opts); + radii = at::empty({nnz, 2}, int_opts); + means2d = at::empty({nnz, 2}, float_opts); + depths = at::empty({nnz}, float_opts); + conics = at::empty({nnz, 3}, float_opts); + if (calc_compensations) { + compensations = at::empty({nnz}, float_opts); + } + + // Second pass: write packed data + AT_DISPATCH_FLOATING_TYPES( + means.scalar_type(), + "projection_ewa_3dgs_packed_fwd_kernel_pass2", + [&] { + d_queue + .parallel_for( + range, + PackedProjectionFwdKernel( + B, + C, + N, + means.data_ptr(), + covars.has_value() ? covars.value().data_ptr() + : nullptr, + quats.has_value() ? quats.value().data_ptr() + : nullptr, + scales.has_value() ? scales.value().data_ptr() + : nullptr, + opacities.has_value() + ? opacities.value().data_ptr() + : nullptr, + viewmats.data_ptr(), + Ks.data_ptr(), + image_width, + image_height, + (scalar_t)eps2d, + (scalar_t)near_plane, + (scalar_t)far_plane, + (scalar_t)radius_clip, + camera_model, + block_accum_inclusive.data_ptr(), + nullptr, // block_cnts + indptr.data_ptr(), + batch_ids.data_ptr(), + camera_ids.data_ptr(), + gaussian_ids.data_ptr(), + radii.data_ptr(), + means2d.data_ptr(), + depths.data_ptr(), + conics.data_ptr(), + calc_compensations ? compensations.data_ptr() + : nullptr + ) + ) + .wait(); + } + ); + + // Set the last element of indptr + if (nrows > 0) { + indptr.index_put_({at::indexing::TensorIndex((int64_t)nrows)}, nnz); + } + + return std::make_tuple( + indptr, + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + compensations + ); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp new file mode 100644 index 00000000..abd70cf7 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp @@ -0,0 +1,66 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/ProjBwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple projection_ewa_simple_bwd( + const at::Tensor means, // [..., C, N, 3] + const at::Tensor covars, // [..., C, N, 3, 3] + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t width, + const uint32_t height, + const CameraModelType camera_model, + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_covars2d // [..., C, N, 2, 2] +) { + DEVICE_GUARD(means); + CHECK_INPUT(means); + CHECK_INPUT2(covars, means); + CHECK_INPUT2(Ks, means); + CHECK_INPUT2(v_means2d, means); + CHECK_INPUT2(v_covars2d, means); + + const uint32_t C = means.size(-3); + const uint32_t N = means.size(-2); + const uint32_t total_gaussians = means.numel() / 3; + + at::Tensor v_means = at::empty_like(means); + at::Tensor v_covars = at::empty_like(covars); + + if (total_gaussians > 0) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = + (total_gaussians + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + ProjBwdKernel kernel( + C, + N, + means.data_ptr(), + covars.data_ptr(), + Ks.data_ptr(), + width, + height, + camera_model, + v_means2d.data_ptr(), + v_covars2d.data_ptr(), + v_means.data_ptr(), + v_covars.data_ptr() + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + } + + return std::make_tuple(v_means, v_covars); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp new file mode 100644 index 00000000..604af3f2 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp @@ -0,0 +1,78 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/ProjFwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple projection_ewa_simple_fwd( + const at::Tensor means, // [C, N, 3] + const at::Tensor covars, // [C, N, 3, 3] + const at::Tensor Ks, // [C, 3, 3] + const uint32_t width, + const uint32_t height, + const CameraModelType camera_model +) { + DEVICE_GUARD(means); + CHECK_INPUT(means); + CHECK_INPUT2(covars, means); + CHECK_INPUT2(Ks, means); + TORCH_CHECK( + means.dim() >= 3, "means must have at least 3 dimensions [..., C, N, 3]" + ); + TORCH_CHECK( + covars.dim() >= 4, + "covars must have at least 4 dimensions [..., C, N, 3, 3]" + ); + TORCH_CHECK( + Ks.dim() >= 3, "Ks must have at least 3 dimensions [..., C, 3, 3]" + ); + + const uint32_t C = means.size(-3); + const uint32_t N = means.size(-2); + const uint32_t total_gaussians = means.numel() / 3; + + auto options = means.options(); + at::DimVector batch_dims(means.sizes().slice(0, means.dim() - 3)); + + at::DimVector means2d_shape = batch_dims; + means2d_shape.insert(means2d_shape.end(), {C, N, 2}); + at::Tensor means2d = at::empty(means2d_shape, options); + + at::DimVector covars2d_shape = batch_dims; + covars2d_shape.insert(covars2d_shape.end(), {C, N, 2, 2}); + at::Tensor covars2d = at::empty(covars2d_shape, covars.options()); + + if (total_gaussians > 0) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = + (total_gaussians + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + ProjFwdKernel kernel( + C, + N, + means.data_ptr(), + covars.data_ptr(), + Ks.data_ptr(), + width, + height, + camera_model, + means2d.data_ptr(), + covars2d.data_ptr() + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + } + + return std::make_tuple(means2d, covars2d); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/projection_ut_3dgs_fused.cpp b/gsplat/sycl/src/projection_ut_3dgs_fused.cpp new file mode 100644 index 00000000..d4cbfd87 --- /dev/null +++ b/gsplat/sycl/src/projection_ut_3dgs_fused.cpp @@ -0,0 +1,45 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +projection_ut_3dgs_fused( + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::optional opacities, // [..., N] optional + const at::Tensor viewmats0, // [..., C, 4, 4] + const at::optional + viewmats1, // [..., C, 4, 4] optional for rolling shutter + const at::Tensor Ks, // [..., C, 3, 3] + const uint32_t image_width, + const uint32_t image_height, + const float eps2d, + const float near_plane, + const float far_plane, + const float radius_clip, + const bool calc_compensations, + const CameraModelType camera_model, + // uncented transform + const UnscentedTransformParameters ut_params, + ShutterType rs_type, + const at::optional + radial_coeffs, // [..., C, 6] or [..., C, 4] optional + const at::optional tangential_coeffs, // [..., C, 2] optional + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs // shared parameters for all cameras +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp new file mode 100644 index 00000000..a774764f --- /dev/null +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -0,0 +1,62 @@ +#include + +#include "Ops.h" +#include "kernels/QuatScaleToCovarPreciBwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple quat_scale_to_covar_preci_bwd( + const at::Tensor quats, // [..., 4] + const at::Tensor scales, // [..., 3] + const bool triu, + const at::optional v_covars, // [..., 3, 3] or [..., 6] + const at::optional v_precis // [..., 3, 3] or [..., 6] +) { + DEVICE_GUARD(quats); + CHECK_INPUT(quats); + CHECK_INPUT2(scales, quats); + if (v_covars.has_value()) { + CHECK_INPUT2(v_covars.value(), quats); + } + if (v_precis.has_value()) { + CHECK_INPUT2(v_precis.value(), quats); + } + TORCH_CHECK( + v_covars.has_value() || v_precis.has_value(), + "Must provide gradients for at least one of covars or precis" + ); + + const int64_t N = quats.numel() / 4; + at::Tensor v_quats = at::empty_like(quats); + at::Tensor v_scales = at::empty_like(scales); + + if (N == 0) { + return std::make_tuple(v_quats, v_scales); + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = (N + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + QuatScaleToCovarPreciBwdKernel kernel( + N, + quats.data_ptr(), + scales.data_ptr(), + v_covars.has_value() ? v_covars.value().data_ptr() : nullptr, + v_precis.has_value() ? v_precis.value().data_ptr() : nullptr, + triu, + v_scales.data_ptr(), + v_quats.data_ptr() + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + + return std::make_tuple(v_quats, v_scales); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp new file mode 100644 index 00000000..c08640ae --- /dev/null +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -0,0 +1,77 @@ +#include + +#include "Ops.h" +#include "kernels/QuatScaleToCovarPreciFwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple quat_scale_to_covar_preci_fwd( + const at::Tensor quats, // [..., 4] + const at::Tensor scales, // [..., 3] + const bool compute_covar, + const bool compute_preci, + const bool triu +) { + DEVICE_GUARD(quats); + CHECK_INPUT(quats); + CHECK_INPUT2(scales, quats); + TORCH_CHECK( + compute_covar || compute_preci, + "Must compute at least one of covar or preci" + ); + + const int64_t N = quats.numel() / 4; + auto options = quats.options(); + + at::Tensor covars; + at::Tensor precis; + + // Create an output shape that preserves the batch dimensions from the input + at::DimVector out_shape(quats.sizes().slice(0, quats.dim() - 1)); + if (triu) { + out_shape.push_back(6); + } else { + out_shape.push_back(3); + out_shape.push_back(3); + } + + if (compute_covar) { + covars = at::empty(out_shape, options); + } else { + covars = at::empty({0}, options); + } + + if (compute_preci) { + precis = at::empty(out_shape, options); + } else { + precis = at::empty({0}, options); + } + + if (N == 0) { + return std::make_tuple(covars, precis); + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = (N + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + QuatScaleToCovarPreciFwdKernel kernel( + N, + quats.data_ptr(), + scales.data_ptr(), + triu, + compute_covar ? covars.data_ptr() : nullptr, + compute_preci ? precis.data_ptr() : nullptr + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + + return std::make_tuple(covars, precis); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp b/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp new file mode 100644 index 00000000..71871778 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp @@ -0,0 +1,28 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple rasterize_to_indices_2dgs( + const uint32_t range_start, + const uint32_t range_end, // iteration steps + const at::Tensor transmittances, // [..., image_height, image_width] + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] + const at::Tensor ray_transforms, // [..., N, 3, 3] + const at::Tensor opacities, // [..., N] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp b/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp new file mode 100644 index 00000000..96726f30 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp @@ -0,0 +1,28 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple rasterize_to_indices_3dgs( + const uint32_t range_start, + const uint32_t range_end, // iteration steps + const at::Tensor transmittances, // [..., image_height, image_width] + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] + const at::Tensor conics, // [..., N, 3] + const at::Tensor opacities, // [..., N] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp new file mode 100644 index 00000000..9858d9e4 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -0,0 +1,301 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/RasterizeToPixels2DGSBwdKernel.hpp" + +namespace gsplat::xpu { + +namespace { + +template +void launch_rasterize_2dgs_bwd_kernel( + // Gaussian parameters + const at::Tensor &means2d, + const at::Tensor &ray_transforms, + const at::Tensor &colors, + const at::Tensor &opacities, + const at::Tensor &normals, + const at::Tensor &densify, + const at::optional &backgrounds, + const at::optional &masks, + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, + // forward outputs + const at::Tensor &render_colors, + const at::Tensor &render_alphas, + const at::Tensor &last_ids, + const at::Tensor &median_ids, + // gradients of outputs + const at::Tensor &v_render_colors, + const at::Tensor &v_render_alphas, + const at::Tensor &v_render_normals, + const at::Tensor &v_render_distort, + const at::Tensor &v_render_median, + // outputs + at::optional v_means2d_abs, + at::Tensor &v_means2d, + at::Tensor &v_ray_transforms, + at::Tensor &v_colors, + at::Tensor &v_opacities, + at::Tensor &v_normals, + at::Tensor &v_densify +) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + bool packed = means2d.dim() == 2; + uint32_t N = packed ? 0 : means2d.size(-2); // number of gaussians + uint32_t I = render_alphas.size(0); // number of images + uint32_t tile_height = tile_offsets.size(-2); + uint32_t tile_width = tile_offsets.size(-1); + uint32_t n_isects = flatten_ids.size(0); + + if (n_isects == 0) { + // Skip kernel launch if there are no intersections + return; + } + + // Define the execution ranges + sycl::range<3> localRange{1, tile_size, tile_size}; + sycl::range<3> globalRange{ + I, tile_height * tile_size, tile_width * tile_size + }; + sycl::nd_range<3> range(globalRange, localRange); + + // Use a fixed chunk size for batching + uint32_t chunk_size = 128; + + auto e = d_queue.submit([&](sycl::handler &cgh) { + // Allocate shared memory + sycl::local_accessor slm_id_batch(chunk_size, cgh); + sycl::local_accessor, 1> slm_xy_opacity( + chunk_size, cgh + ); + sycl::local_accessor, 1> slm_u_Ms(chunk_size, cgh); + sycl::local_accessor, 1> slm_v_Ms(chunk_size, cgh); + sycl::local_accessor, 1> slm_w_Ms(chunk_size, cgh); + sycl::local_accessor, 1> slm_rgbs( + chunk_size, cgh + ); + sycl::local_accessor, 1> slm_normals( + chunk_size, cgh + ); + + RasterizeToPixels2DGSBwdKernel kernel( + I, + N, + n_isects, + packed, + chunk_size, + reinterpret_cast *>( + means2d.data_ptr() + ), + ray_transforms.data_ptr(), + colors.data_ptr(), + opacities.data_ptr(), + normals.data_ptr(), + backgrounds.has_value() ? backgrounds.value().data_ptr() + : nullptr, + masks.has_value() ? masks.value().data_ptr() : nullptr, + image_width, + image_height, + tile_size, + tile_width, + tile_height, + tile_offsets.data_ptr(), + flatten_ids.data_ptr(), + render_colors.data_ptr(), + render_alphas.data_ptr(), + last_ids.data_ptr(), + median_ids.data_ptr(), + v_render_colors.data_ptr(), + v_render_alphas.data_ptr(), + v_render_normals.data_ptr(), + v_render_distort.data_ptr(), + v_render_median.data_ptr(), + v_means2d_abs.has_value() + ? reinterpret_cast *>( + v_means2d_abs.value().data_ptr() + ) + : nullptr, + reinterpret_cast *>(v_means2d.data_ptr() + ), + v_ray_transforms.data_ptr(), + v_colors.data_ptr(), + v_opacities.data_ptr(), + v_normals.data_ptr(), + v_densify.data_ptr(), + slm_id_batch, + slm_xy_opacity, + slm_u_Ms, + slm_v_Ms, + slm_w_Ms, + slm_rgbs, + slm_normals + ); + + cgh.parallel_for(range, kernel); + }); + e.wait(); +} + +} // anonymous namespace + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +rasterize_to_pixels_2dgs_bwd( + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] or [nnz, 2] + const at::Tensor ray_transforms, // [..., N, 3, 3] or [nnz, 3, 3] + const at::Tensor colors, // [..., N, channels] or [nnz, channels] + const at::Tensor opacities, // [..., N] or [nnz] + const at::Tensor normals, // [..., N, 3] or [nnz, 3] + const at::Tensor densify, // [..., N, 2] or [nnz, 2] + const at::optional backgrounds, // [..., channels] + const at::optional masks, // [..., tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids, // [n_isects] + // forward outputs + const at::Tensor + render_colors, // [..., image_height, image_width, channels] + const at::Tensor render_alphas, // [..., image_height, image_width] + const at::Tensor last_ids, // [..., image_height, image_width] + const at::Tensor median_ids, // [..., image_height, image_width] + // gradients of outputs + const at::Tensor + v_render_colors, // [..., image_height, image_width, channels] + const at::Tensor v_render_alphas, // [..., image_height, image_width] + const at::Tensor v_render_normals, // [..., image_height, image_width, 3] + const at::Tensor v_render_distort, // [..., image_height, image_width] + const at::Tensor v_render_median, // [..., image_height, image_width] + bool absgrad +) { + DEVICE_GUARD(means2d); + // Check input tensors are contiguous and on the same device + CHECK_INPUT(means2d); + CHECK_INPUT2(ray_transforms, means2d); + CHECK_INPUT2(colors, means2d); + CHECK_INPUT2(opacities, means2d); + CHECK_INPUT2(normals, means2d); + if (backgrounds.has_value()) { + CHECK_INPUT2(backgrounds.value(), means2d); + } + if (masks.has_value()) { + CHECK_INPUT2(masks.value(), means2d); + } + CHECK_INPUT2(tile_offsets, means2d); + CHECK_INPUT2(flatten_ids, means2d); + CHECK_INPUT2(render_colors, means2d); + CHECK_INPUT2(render_alphas, means2d); + CHECK_INPUT2(last_ids, means2d); + CHECK_INPUT2(median_ids, means2d); + CHECK_INPUT2(v_render_colors, means2d); + CHECK_INPUT2(v_render_alphas, means2d); + CHECK_INPUT2(v_render_normals, means2d); + CHECK_INPUT2(v_render_distort, means2d); + CHECK_INPUT2(v_render_median, means2d); + + uint32_t channels = colors.size(-1); + + // Create output tensors + auto options = means2d.options().dtype(torch::kFloat32); + at::Tensor v_means2d_abs; + if (absgrad) { + v_means2d_abs = at::zeros_like(means2d, options); + } + at::Tensor v_means2d = at::zeros_like(means2d, options); + at::Tensor v_ray_transforms = at::zeros_like(ray_transforms, options); + at::Tensor v_colors = at::zeros_like(colors, options); + at::Tensor v_opacities = at::zeros_like(opacities, options); + at::Tensor v_normals = at::zeros_like(normals, options); + at::Tensor v_densify = at::zeros_like(densify, options); + + // Launch kernel with appropriate dimension +#define __GS__CALL_(DIM) \ + case DIM: \ + launch_rasterize_2dgs_bwd_kernel( \ + means2d, \ + ray_transforms, \ + colors, \ + opacities, \ + normals, \ + densify, \ + backgrounds, \ + masks, \ + image_width, \ + image_height, \ + tile_size, \ + tile_offsets, \ + flatten_ids, \ + render_colors, \ + render_alphas, \ + last_ids, \ + median_ids, \ + v_render_colors, \ + v_render_alphas, \ + v_render_normals, \ + v_render_distort, \ + v_render_median, \ + absgrad ? c10::optional(v_means2d_abs) : c10::nullopt, \ + v_means2d, \ + v_ray_transforms, \ + v_colors, \ + v_opacities, \ + v_normals, \ + v_densify \ + ); \ + break; + + switch (channels) { + __GS__CALL_(1); + __GS__CALL_(2); + __GS__CALL_(3); + __GS__CALL_(4); + __GS__CALL_(5); + __GS__CALL_(8); + __GS__CALL_(9); + __GS__CALL_(16); + __GS__CALL_(17); + __GS__CALL_(32); + __GS__CALL_(33); + __GS__CALL_(64); + __GS__CALL_(65); + __GS__CALL_(128); + __GS__CALL_(129); + __GS__CALL_(256); + __GS__CALL_(257); + __GS__CALL_(512); + __GS__CALL_(513); + default: + TORCH_CHECK(false, "Unsupported number of channels: ", channels); + } +#undef __GS__CALL_ + + return std::make_tuple( + v_means2d_abs, + v_means2d, + v_ray_transforms, + v_colors, + v_opacities, + v_normals, + v_densify + ); +} + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp new file mode 100644 index 00000000..0cfa20e3 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp @@ -0,0 +1,264 @@ +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/RasterizeToPixels2DGSFwdKernel.hpp" + +namespace gsplat::xpu { + +namespace { + +template +void launch_rasterize_2dgs_kernel( + // Gaussian parameters + const at::Tensor &means2d, + const at::Tensor &ray_transforms, + const at::Tensor &colors, + const at::Tensor &opacities, + const at::Tensor &normals, + const at::optional &backgrounds, + const at::optional &masks, + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, + // other params + bool packed, + uint32_t I, + uint32_t N, + uint32_t tile_height, + uint32_t tile_width, + uint32_t n_isects, + // outputs + at::Tensor &renders, + at::Tensor &alphas, + at::Tensor &render_normals, + at::Tensor &render_distort, + at::Tensor &render_median, + at::Tensor &last_ids, + at::Tensor &median_ids +) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + // Define the execution ranges + sycl::range<3> localRange{1, tile_size, tile_size}; + sycl::range<3> globalRange{ + I, tile_height * tile_size, tile_width * tile_size + }; + sycl::nd_range<3> range(globalRange, localRange); + + // Use a fixed chunk size for batching - don't make it constexpr with + // tile_size + uint32_t chunk_size = + 128; // Fixed size that's similar to what would be used + + auto e = d_queue.submit([&](sycl::handler &cgh) { + // Allocate shared memory + sycl::local_accessor slm_id_batch(chunk_size, cgh); + sycl::local_accessor, 1> slm_xy_opacity( + chunk_size, cgh + ); + sycl::local_accessor, 1> slm_u_Ms(chunk_size, cgh); + sycl::local_accessor, 1> slm_v_Ms(chunk_size, cgh); + sycl::local_accessor, 1> slm_w_Ms(chunk_size, cgh); + + RasterizeToPixels2DGSFwdKernel kernel( + I, + N, + n_isects, + packed, + chunk_size, + reinterpret_cast *>( + means2d.data_ptr() + ), + ray_transforms.data_ptr(), + colors.data_ptr(), + opacities.data_ptr(), + normals.data_ptr(), + backgrounds.has_value() ? backgrounds.value().data_ptr() + : nullptr, + masks.has_value() ? masks.value().data_ptr() : nullptr, + image_width, + image_height, + tile_size, + tile_width, + tile_height, + tile_offsets.data_ptr(), + flatten_ids.data_ptr(), + renders.data_ptr(), + alphas.data_ptr(), + render_normals.data_ptr(), + render_distort.data_ptr(), + render_median.data_ptr(), + last_ids.data_ptr(), + median_ids.data_ptr(), + slm_id_batch, + slm_xy_opacity, + slm_u_Ms, + slm_v_Ms, + slm_w_Ms + ); + + cgh.parallel_for(range, kernel); + }); + e.wait(); +} + +} // anonymous namespace + +std::tuple< + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor, + at::Tensor> +rasterize_to_pixels_2dgs_fwd( + // Gaussian parameters + const at::Tensor means2d, // [..., N, 2] or [nnz, 2] + const at::Tensor ray_transforms, // [..., N, 3, 3] or [nnz, 3, 3] + const at::Tensor colors, // [..., N, channels] or [nnz, channels] + const at::Tensor opacities, // [..., N] or [nnz] + const at::Tensor normals, // [..., N, 3] or [nnz, 3] + const at::optional backgrounds, // [..., channels] + const at::optional masks, // [..., tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, // [..., tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +) { + DEVICE_GUARD(means2d); + // Check input tensors are contiguous and on the same device + CHECK_INPUT(means2d); + CHECK_INPUT2(ray_transforms, means2d); + CHECK_INPUT2(colors, means2d); + CHECK_INPUT2(opacities, means2d); + CHECK_INPUT2(normals, means2d); + CHECK_INPUT2(tile_offsets, means2d); + CHECK_INPUT2(flatten_ids, means2d); + if (backgrounds.has_value()) + CHECK_INPUT2(backgrounds.value(), means2d); + if (masks.has_value()) + CHECK_INPUT2(masks.value(), means2d); + + // Get dimensions + bool packed = means2d.dim() == 2; + uint32_t N = packed ? 0 : means2d.size(-2); // number of gaussians + at::DimVector image_dims( + tile_offsets.sizes().slice(0, tile_offsets.dim() - 2) + ); + uint32_t tile_height = tile_offsets.size(-2); + uint32_t tile_width = tile_offsets.size(-1); + uint32_t I = + tile_offsets.numel() / (tile_height * tile_width); // number of images + uint32_t n_isects = flatten_ids.size(0); // number of intersections + uint32_t channels = colors.size(-1); // color dimension + + // Create output tensors + auto options_float = means2d.options().dtype(torch::kFloat32); + auto options_int = means2d.options().dtype(torch::kInt32); + + at::DimVector renders_dims(image_dims); + renders_dims.append({image_height, image_width, channels}); + at::Tensor renders = at::zeros(renders_dims, options_float); + + at::DimVector alphas_dims(image_dims); + alphas_dims.append({image_height, image_width, 1}); + at::Tensor alphas = at::zeros(alphas_dims, options_float); + + at::DimVector render_normals_dims(image_dims); + render_normals_dims.append({image_height, image_width, 3}); + at::Tensor render_normals = at::zeros(render_normals_dims, options_float); + + at::DimVector render_distort_dims(image_dims); + render_distort_dims.append({image_height, image_width, 1}); + at::Tensor render_distort = at::zeros(render_distort_dims, options_float); + + at::DimVector render_median_dims(image_dims); + render_median_dims.append({image_height, image_width, 1}); + at::Tensor render_median = at::zeros(render_median_dims, options_float); + + at::DimVector last_ids_dims(image_dims); + last_ids_dims.append({image_height, image_width}); + at::Tensor last_ids = at::zeros(last_ids_dims, options_int); + + at::DimVector median_ids_dims(image_dims); + median_ids_dims.append({image_height, image_width}); + at::Tensor median_ids = at::zeros(median_ids_dims, options_int); + + // Launch kernel with appropriate dimension +#define __GS__CALL_(DIM) \ + case DIM: \ + launch_rasterize_2dgs_kernel( \ + means2d, \ + ray_transforms, \ + colors, \ + opacities, \ + normals, \ + backgrounds, \ + masks, \ + image_width, \ + image_height, \ + tile_size, \ + tile_offsets, \ + flatten_ids, \ + packed, \ + I, \ + N, \ + tile_height, \ + tile_width, \ + n_isects, \ + renders, \ + alphas, \ + render_normals, \ + render_distort, \ + render_median, \ + last_ids, \ + median_ids \ + ); \ + break; + + switch (channels) { + __GS__CALL_(1); + __GS__CALL_(2); + __GS__CALL_(3); + __GS__CALL_(4); + __GS__CALL_(5); + __GS__CALL_(8); + __GS__CALL_(9); + __GS__CALL_(16); + __GS__CALL_(17); + __GS__CALL_(32); + __GS__CALL_(33); + __GS__CALL_(64); + __GS__CALL_(65); + __GS__CALL_(128); + __GS__CALL_(129); + __GS__CALL_(256); + __GS__CALL_(257); + __GS__CALL_(512); + __GS__CALL_(513); + default: + TORCH_CHECK(false, "Unsupported number of channels: ", channels); + } +#undef __GS__CALL_ + + return std::make_tuple( + renders, + alphas, + render_normals, + render_distort, + render_median, + last_ids, + median_ids + ); +} + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp new file mode 100644 index 00000000..629bafa9 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp @@ -0,0 +1,255 @@ +#include + +#include "Ops.h" +#include "kernels/RasterizeToPixelsBwdKernel.hpp" + +namespace gsplat::xpu { + +namespace { + +template +void launch_rasterize_bwd_kernel( + // Gaussian parameters + const at::Tensor &means2d, + const at::Tensor &conics, + const at::Tensor &colors, + const at::Tensor &opacities, + const at::optional &backgrounds, + const at::optional &masks, + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, + // forward outputs + const at::Tensor &render_alphas, + const at::Tensor &last_ids, + // gradients of outputs + const at::Tensor &v_render_colors, + const at::Tensor &v_render_alphas, + // options and derived params + bool absgrad, + bool packed, + uint32_t C, + uint32_t N, + uint32_t n_isects, + uint32_t tile_height, + uint32_t tile_width, + // output grads + at::Tensor &v_means2d, + at::Tensor &v_conics, + at::Tensor &v_colors, + at::Tensor &v_opacities, + at::Tensor &v_means2d_abs +) { + if (n_isects == 0) { + return; + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + sycl::range<3> localRange{1, tile_size, tile_size}; + sycl::range<3> globalRange{ + C, tile_height * tile_size, tile_width * tile_size + }; + sycl::nd_range<3> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + constexpr uint32_t CHUNK_SIZE = 256; + sycl::range<1> slm_range(CHUNK_SIZE); + + sycl::local_accessor slm_flatten_ids(slm_range, cgh); + sycl::local_accessor, 1> slm_means2d( + slm_range, cgh + ); + sycl::local_accessor slm_opacities(slm_range, cgh); + sycl::local_accessor, 1> slm_conics( + slm_range, cgh + ); + sycl::local_accessor, 1> slm_color; + if constexpr (BufferType::isVec && COLOR_DIM <= 4) { + slm_color = + sycl::local_accessor, 1>( + slm_range, cgh + ); + } + + RasterizeToPixelsBwdKernel kernel( + C, + N, + n_isects, + packed, + 0, + nullptr, // concat_stride, concatenated_data + reinterpret_cast *>( + means2d.data_ptr() + ), + reinterpret_cast *>(conics.data_ptr()), + colors.data_ptr(), + opacities.data_ptr(), + backgrounds.has_value() ? backgrounds.value().data_ptr() + : nullptr, + masks.has_value() ? masks.value().data_ptr() : nullptr, + image_width, + image_height, + tile_size, + tile_width, + tile_height, + tile_offsets.data_ptr(), + flatten_ids.data_ptr(), + render_alphas.data_ptr(), + last_ids.data_ptr(), + v_render_colors.data_ptr(), + v_render_alphas.data_ptr(), + absgrad ? reinterpret_cast *>( + v_means2d_abs.data_ptr() + ) + : nullptr, + reinterpret_cast *>(v_means2d.data_ptr() + ), + reinterpret_cast *>(v_conics.data_ptr()), + v_colors.data_ptr(), + v_opacities.data_ptr(), + slm_flatten_ids, + slm_means2d, + slm_opacities, + slm_conics, + slm_color + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); +} + +} // anonymous namespace + +std::tuple +rasterize_to_pixels_3dgs_bwd( + // Gaussian parameters + const at::Tensor means2d, // [..., C, N, 2] or [C, N, 2] + const at::Tensor conics, // [..., C, N, 3] or [C, N, 3] + const at::Tensor colors, // [..., C, N, COLOR_DIM] or [C, N, COLOR_DIM] + const at::Tensor opacities, // [..., C, N] or [C, N] + const at::optional + backgrounds, // [..., C, COLOR_DIM] or [C, COLOR_DIM] optional + const at::optional + masks, // [..., C, image_height, image_width] optional + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, + const at::Tensor flatten_ids, + // forward outputs + const at::Tensor render_alphas, // [..., C, image_height, image_width, 1] + const at::Tensor last_ids, // [..., C, image_height, image_width] + // gradients of outputs + const at::Tensor + v_render_colors, // [..., C, image_height, image_width, COLOR_DIM] + const at::Tensor v_render_alphas, // [..., C, image_height, image_width, 1] + // options + bool absgrad +) { + DEVICE_GUARD(means2d); + // Check input tensors are contiguous and on the same device + CHECK_INPUT(means2d); + CHECK_INPUT2(conics, means2d); + CHECK_INPUT2(colors, means2d); + CHECK_INPUT2(opacities, means2d); + CHECK_INPUT2(tile_offsets, means2d); + CHECK_INPUT2(flatten_ids, means2d); + CHECK_INPUT2(render_alphas, means2d); + CHECK_INPUT2(last_ids, means2d); + CHECK_INPUT2(v_render_colors, means2d); + CHECK_INPUT2(v_render_alphas, means2d); + if (backgrounds.has_value()) + CHECK_INPUT2(backgrounds.value(), means2d); + if (masks.has_value()) + CHECK_INPUT2(masks.value(), means2d); + + TORCH_CHECK(means2d.dim() >= 2, "means2d must have at least 2 dimensions"); + TORCH_CHECK(colors.dim() >= 2, "colors must have at least 2 dimensions"); + + // --- Parameter Derivation --- + const uint32_t COLOR_DIM = colors.size(-1); + const bool packed = means2d.dim() == 2; + const uint32_t C = tile_offsets.size(0); + const uint32_t N = packed ? 0 : means2d.size(-2); + const uint32_t n_isects = flatten_ids.size(0); + const uint32_t tile_height = tile_offsets.size(1); + const uint32_t tile_width = tile_offsets.size(2); + + at::Tensor v_means2d = at::zeros_like(means2d); + at::Tensor v_conics = at::zeros_like(conics); + at::Tensor v_colors = at::zeros_like(colors); + at::Tensor v_opacities = at::zeros_like(opacities); + at::Tensor v_means2d_abs = + absgrad ? at::zeros_like(means2d) : at::empty({0}, means2d.options()); + +#define __GS_BWD_CALL_(DIM) \ + case DIM: \ + launch_rasterize_bwd_kernel( \ + means2d, \ + conics, \ + colors, \ + opacities, \ + backgrounds, \ + masks, \ + image_width, \ + image_height, \ + tile_size, \ + tile_offsets, \ + flatten_ids, \ + render_alphas, \ + last_ids, \ + v_render_colors, \ + v_render_alphas, \ + absgrad, \ + packed, \ + C, \ + N, \ + n_isects, \ + tile_height, \ + tile_width, \ + v_means2d, \ + v_conics, \ + v_colors, \ + v_opacities, \ + v_means2d_abs \ + ); \ + break; + + switch (COLOR_DIM) { + __GS_BWD_CALL_(1); + __GS_BWD_CALL_(2); + __GS_BWD_CALL_(3); + __GS_BWD_CALL_(4); + __GS_BWD_CALL_(5); + __GS_BWD_CALL_(8); + __GS_BWD_CALL_(9); + __GS_BWD_CALL_(16); + __GS_BWD_CALL_(17); + __GS_BWD_CALL_(32); + __GS_BWD_CALL_(33); + __GS_BWD_CALL_(64); + __GS_BWD_CALL_(65); + __GS_BWD_CALL_(128); + __GS_BWD_CALL_(129); + __GS_BWD_CALL_(256); + __GS_BWD_CALL_(257); + __GS_BWD_CALL_(512); + __GS_BWD_CALL_(513); + default: + TORCH_CHECK(false, "Unsupported number of channels: ", COLOR_DIM); + } +#undef __GS_BWD_CALL_ + + return std::make_tuple( + v_means2d_abs, v_means2d, v_conics, v_colors, v_opacities + ); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp new file mode 100644 index 00000000..802adfc5 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp @@ -0,0 +1,212 @@ +#include + +#include "Ops.h" +#include "kernels/RasterizeToPixelsFwdKernel.hpp" + +namespace gsplat::xpu { + +namespace { + +template +void launch_rasterize_kernel( + // Gaussian parameters + const at::Tensor &means2d, + const at::Tensor &conics, + const at::Tensor &colors, + const at::Tensor &opacities, + const at::optional &backgrounds, + const at::optional &masks, + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, + // other params + bool packed, + uint32_t C, + uint32_t N, + uint32_t tile_height, + uint32_t tile_width, + // outputs + at::Tensor &renders, + at::Tensor &alphas, + at::Tensor &last_ids +) { + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + sycl::range<3> localRange{1, tile_size, tile_size}; + sycl::range<3> globalRange{ + C, tile_height * tile_size, tile_width * tile_size + }; + sycl::nd_range<3> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + constexpr uint32_t CHUNK_SIZE = 128; + sycl::range<1> slm_range(tile_size * tile_size); + + sycl::local_accessor slm_flatten_ids(slm_range, cgh); + sycl::local_accessor, 1> slm_means2d( + slm_range, cgh + ); + sycl::local_accessor slm_opacities(slm_range, cgh); + sycl::local_accessor, 1> slm_conics(slm_range, cgh); + sycl::local_accessor, 1> slm_color; + if constexpr (BufferType::isVec && COLOR_DIM <= 4) { + slm_color = sycl::local_accessor, 1>( + slm_range, cgh + ); + } + + RasterizeToPixelsFwdKernel kernel( + C, + N, + flatten_ids.size(0), + packed, + 0, + nullptr, // concat_stride, concatenated_data + reinterpret_cast *>( + means2d.data_ptr() + ), + reinterpret_cast *>(conics.data_ptr()), + colors.data_ptr(), + opacities.data_ptr(), + backgrounds.has_value() ? backgrounds.value().data_ptr() + : nullptr, + masks.has_value() ? masks.value().data_ptr() : nullptr, + image_width, + image_height, + tile_size, + tile_width, + tile_height, + tile_offsets.data_ptr(), + flatten_ids.data_ptr(), + renders.data_ptr(), + alphas.data_ptr(), + last_ids.data_ptr(), + slm_flatten_ids, + slm_means2d, + slm_opacities, + slm_conics, + slm_color + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); +} +} // anonymous namespace + +std::tuple rasterize_to_pixels_3dgs_fwd( + // Gaussian parameters + const at::Tensor means2d, // [..., C, N, 2] or [C, N, 2] + const at::Tensor conics, // [..., C, N, 3] or [C, N, 3] + const at::Tensor colors, // [..., C, N, COLOR_DIM] or [C, N, COLOR_DIM] + const at::Tensor opacities, // [..., C, N] or [C, N] + const at::optional + backgrounds, // [..., C, COLOR_DIM] or [C, COLOR_DIM] optional + const at::optional + masks, // [..., C, image_height, image_width] optional + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // intersections + const at::Tensor tile_offsets, + const at::Tensor flatten_ids +) { + DEVICE_GUARD(means2d); + CHECK_INPUT(means2d); + CHECK_INPUT2(conics, means2d); + CHECK_INPUT2(colors, means2d); + CHECK_INPUT2(opacities, means2d); + CHECK_INPUT2(tile_offsets, means2d); + CHECK_INPUT2(flatten_ids, means2d); + if (backgrounds.has_value()) + CHECK_INPUT2(backgrounds.value(), means2d); + if (masks.has_value()) + CHECK_INPUT2(masks.value(), means2d); + + TORCH_CHECK(means2d.dim() >= 2, "means2d must have at least 2 dimensions"); + TORCH_CHECK(colors.dim() >= 2, "colors must have at least 2 dimensions"); + + const uint32_t channels = colors.size(-1); + const bool packed = means2d.dim() == 2; + const uint32_t C = tile_offsets.size(0); + const uint32_t N = packed ? 0 : means2d.size(-2); + const uint32_t tile_height = tile_offsets.size(1); + const uint32_t tile_width = tile_offsets.size(2); + + auto options_float = means2d.options().dtype(torch::kFloat32); + auto options_int = means2d.options().dtype(torch::kInt32); + at::DimVector image_dims( + tile_offsets.sizes().slice(0, tile_offsets.dim() - 2) + ); + + at::DimVector out_shape_renders = image_dims; + out_shape_renders.append({image_height, image_width, channels}); + + at::DimVector out_shape_alphas = image_dims; + out_shape_alphas.append({image_height, image_width, 1}); + + at::DimVector out_shape_last_ids = image_dims; + out_shape_last_ids.append({image_height, image_width}); + + at::Tensor renders = at::empty(out_shape_renders, options_float); + at::Tensor alphas = at::empty(out_shape_alphas, options_float); + at::Tensor last_ids = at::empty(out_shape_last_ids, options_int); + +#define __GS__CALL_(DIM) \ + case DIM: \ + launch_rasterize_kernel( \ + means2d, \ + conics, \ + colors, \ + opacities, \ + backgrounds, \ + masks, \ + image_width, \ + image_height, \ + tile_size, \ + tile_offsets, \ + flatten_ids, \ + packed, \ + C, \ + N, \ + tile_height, \ + tile_width, \ + renders, \ + alphas, \ + last_ids \ + ); \ + break; + + switch (channels) { + __GS__CALL_(1); + __GS__CALL_(2); + __GS__CALL_(3); + __GS__CALL_(4); + __GS__CALL_(5); + __GS__CALL_(8); + __GS__CALL_(9); + __GS__CALL_(16); + __GS__CALL_(17); + __GS__CALL_(32); + __GS__CALL_(33); + __GS__CALL_(64); + __GS__CALL_(65); + __GS__CALL_(128); + __GS__CALL_(129); + __GS__CALL_(256); + __GS__CALL_(257); + __GS__CALL_(512); + __GS__CALL_(513); + default: + TORCH_CHECK(false, "Unsupported number of channels: ", channels); + } +#undef __GS__CALL_ + + return std::make_tuple(renders, alphas, last_ids); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp new file mode 100644 index 00000000..e9233c5c --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp @@ -0,0 +1,51 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple +rasterize_to_pixels_from_world_3dgs_bwd( + // Gaussian parameters + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor colors, // [..., C, N, 3] or [nnz, 3] + const at::Tensor opacities, // [..., C, N] or [nnz] + const at::optional backgrounds, // [..., C, 3] + const at::optional masks, // [..., C, tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // camera + const at::Tensor viewmats0, // [..., C, 4, 4] + const at::optional + viewmats1, // [..., C, 4, 4] optional for rolling shutter + const at::Tensor Ks, // [..., C, 3, 3] + const CameraModelType camera_model, + // uncented transform + const UnscentedTransformParameters ut_params, + ShutterType rs_type, + const at::optional + radial_coeffs, // [..., C, 6] or [..., C, 4] optional + const at::optional tangential_coeffs, // [..., C, 2] optional + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs, // shared parameters for all cameras + // intersections + const at::Tensor tile_offsets, // [..., C, tile_height, tile_width] + const at::Tensor flatten_ids, // [n_isects] + // forward outputs + const at::Tensor render_alphas, // [..., C, image_height, image_width, 1] + const at::Tensor last_ids, // [..., C, image_height, image_width] + // gradients of outputs + const at::Tensor v_render_colors, // [..., C, image_height, image_width, 3] + const at::Tensor v_render_alphas // [..., C, image_height, image_width, 1] +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp new file mode 100644 index 00000000..e0bb0d64 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp @@ -0,0 +1,45 @@ + +#include + +#include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { + +std::tuple +rasterize_to_pixels_from_world_3dgs_fwd( + // Gaussian parameters + const at::Tensor means, // [..., N, 3] + const at::Tensor quats, // [..., N, 4] + const at::Tensor scales, // [..., N, 3] + const at::Tensor colors, // [..., C, N, channels] or [nnz, channels] + const at::Tensor opacities, // [..., C, N] or [nnz] + const at::optional backgrounds, // [..., C, channels] + const at::optional masks, // [..., C, tile_height, tile_width] + // image size + const uint32_t image_width, + const uint32_t image_height, + const uint32_t tile_size, + // camera + const at::Tensor viewmats0, // [..., C, 4, 4] + const at::optional + viewmats1, // [..., C, 4, 4] optional for rolling shutter + const at::Tensor Ks, // [..., C, 3, 3] + const CameraModelType camera_model, + // uncented transform + const UnscentedTransformParameters ut_params, + ShutterType rs_type, + const at::optional + radial_coeffs, // [..., C, 6] or [..., C, 4] optional + const at::optional tangential_coeffs, // [..., C, 2] optional + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs, // shared parameters for all cameras + // intersections + const at::Tensor tile_offsets, // [..., C, tile_height, tile_width] + const at::Tensor flatten_ids // [n_isects] +) { + throw std::runtime_error(std::string(__func__) + " is not implemented"); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/relocation.cpp b/gsplat/sycl/src/relocation.cpp new file mode 100644 index 00000000..17777a7d --- /dev/null +++ b/gsplat/sycl/src/relocation.cpp @@ -0,0 +1,53 @@ + +#include + +#include "Common.h" +#include "Ops.h" +#include "kernels/RelocationKernel.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { + +std::tuple relocation( + at::Tensor opacities, // [N] + at::Tensor scales, // [N, 3] + at::Tensor ratios, // [N] + at::Tensor binoms, // [n_max, n_max] + const int n_max +) { + DEVICE_GUARD(opacities); + // Check input tensors are contiguous and on the same device + CHECK_INPUT(opacities); + CHECK_INPUT2(scales, opacities); + CHECK_INPUT2(ratios, opacities); + CHECK_INPUT2(binoms, opacities); + if (opacities.size(0) == 0) { + return std::make_tuple( + at::empty_like(opacities), at::empty_like(scales) + ); + } + at::Tensor new_opacities = at::empty_like(opacities); + at::Tensor new_scales = at::empty_like(scales); + + AT_DISPATCH_FLOATING_TYPES(opacities.scalar_type(), "relocation", ([&] { + auto &d_queue = + at::xpu::getCurrentXPUStream().queue(); + auto e = d_queue.parallel_for( + sycl::range<1>(opacities.size(0)), + kernels::RelocationKernel( + opacities.data_ptr(), + scales.data_ptr(), + ratios.data_ptr(), + binoms.data_ptr(), + n_max, + new_opacities.data_ptr(), + new_scales.data_ptr() + ) + ); + e.wait(); + })); + + return std::make_tuple(new_opacities, new_scales); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/spherical_harmonics_bwd.cpp b/gsplat/sycl/src/spherical_harmonics_bwd.cpp new file mode 100644 index 00000000..acb8fd1f --- /dev/null +++ b/gsplat/sycl/src/spherical_harmonics_bwd.cpp @@ -0,0 +1,65 @@ +#include + +#include "Ops.h" +#include "kernels/ComputeShBwdKernel.hpp" + +namespace gsplat::xpu { + +std::tuple spherical_harmonics_bwd( + const uint32_t K, + const uint32_t degrees_to_use, + const at::Tensor dirs, // [..., 3] + const at::Tensor coeffs, // [..., K, 3] + const at::optional masks, // [...] + const at::Tensor v_colors, // [..., 3] + bool compute_v_dirs +) { + DEVICE_GUARD(dirs); + CHECK_INPUT(dirs); + CHECK_INPUT2(coeffs, dirs); + CHECK_INPUT2(v_colors, dirs); + if (masks.has_value()) { + CHECK_INPUT2(masks.value(), dirs); + } + + TORCH_CHECK(v_colors.size(-1) == 3, "v_colors must have last dimension 3"); + TORCH_CHECK(coeffs.size(-1) == 3, "coeffs must have last dimension 3"); + TORCH_CHECK(dirs.size(-1) == 3, "dirs must have last dimension 3"); + + const uint32_t N = dirs.numel() / 3; + + at::Tensor v_coeffs = at::zeros_like(coeffs); + at::Tensor v_dirs = + compute_v_dirs ? at::zeros_like(dirs) : at::empty({0}, dirs.options()); + + if (N == 0) { + return std::make_tuple(v_coeffs, v_dirs); + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = (N * 3 + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + ComputeShBwdKernel kernel( + N, + K, + degrees_to_use, + reinterpret_cast *>(dirs.data_ptr()), + coeffs.data_ptr(), + masks.has_value() ? masks.value().data_ptr() : nullptr, + v_colors.data_ptr(), + v_coeffs.data_ptr(), + compute_v_dirs ? v_dirs.data_ptr() : nullptr + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + + return std::make_tuple(v_coeffs, v_dirs); +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/spherical_harmonics_fwd.cpp b/gsplat/sycl/src/spherical_harmonics_fwd.cpp new file mode 100644 index 00000000..3885af3f --- /dev/null +++ b/gsplat/sycl/src/spherical_harmonics_fwd.cpp @@ -0,0 +1,62 @@ +#include + +#include "Ops.h" +#include "kernels/ComputeShFwdKernel.hpp" + +namespace gsplat::xpu { + +at::Tensor spherical_harmonics_fwd( + const uint32_t degrees_to_use, + const at::Tensor dirs, // [..., 3] + const at::Tensor coeffs, // [..., K, 3] + const at::optional masks // [...] +) { + DEVICE_GUARD(dirs); + CHECK_INPUT(dirs); + CHECK_INPUT2(coeffs, dirs); + if (masks.has_value()) { + CHECK_INPUT2(masks.value(), dirs); + } + TORCH_CHECK( + dirs.size(-1) == 3, + "Input 'dirs' tensor must have the last dimension of size 3." + ); + TORCH_CHECK( + coeffs.size(-1) == 3, + "Input 'coeffs' tensor must have the last dimension of size 3." + ); + + const uint32_t K = coeffs.size(-2); + const uint32_t N = dirs.numel() / 3; + + at::Tensor colors = at::empty_like(dirs); + + if (N == 0) { + return colors; + } + + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); + + size_t numWorkGrps = (N * 3 + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + sycl::range<1> localRange(GSPLAT_N_THREADS); + sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); + sycl::nd_range<1> range(globalRange, localRange); + + auto e = d_queue.submit([&](sycl::handler &cgh) { + ComputeShFwdKernel kernel( + N, + K, + degrees_to_use, + reinterpret_cast *>(dirs.data_ptr()), + coeffs.data_ptr(), + masks.has_value() ? masks.value().data_ptr() : nullptr, + colors.data_ptr() + ); + cgh.parallel_for(range, kernel); + }); + e.wait(); + + return colors; +} + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/utils.py b/gsplat/utils.py index 4924091f..d527e827 100644 --- a/gsplat/utils.py +++ b/gsplat/utils.py @@ -231,7 +231,7 @@ def depth_to_normal( return normals -def get_projection_matrix(znear, zfar, fovX, fovY, device="cuda"): +def get_projection_matrix(znear, zfar, fovX, fovY, device=None) -> Tensor: """Create OpenGL-style projection matrix""" tanHalfFovY = math.tan((fovY / 2)) tanHalfFovX = math.tan((fovX / 2)) @@ -241,6 +241,12 @@ def get_projection_matrix(znear, zfar, fovX, fovY, device="cuda"): right = tanHalfFovX * znear left = -right + if device is None: + device = ( + torch.accelerator.current_accelerator() + if torch.accelerator.is_available() + else torch.device("cpu") + ) P = torch.zeros(4, 4, device=device) z_sign = 1.0 diff --git a/profiling/batch.py b/profiling/batch.py index 6aaae54a..67b51840 100644 --- a/profiling/batch.py +++ b/profiling/batch.py @@ -11,6 +11,7 @@ import torch from typing_extensions import Callable, Literal +from gsplat import torch_acc, BACKEND from gsplat._helper import load_test_data from gsplat.distributed import cli from gsplat.rendering import rasterization @@ -22,17 +23,17 @@ "4k": (3840, 2160), } -device = torch.device("cuda") +device = torch_acc._device(0) def timeit(repeats: int, f: Callable, *args, **kwargs) -> float: for _ in range(5): # warmup f(*args, **kwargs) - torch.cuda.synchronize() + torch_acc.synchronize() start = time.time() for _ in range(repeats): results = f(*args, **kwargs) - torch.cuda.synchronize() + torch_acc.synchronize() end = time.time() return (end - start) / repeats, results @@ -79,8 +80,8 @@ def main( Ks[..., 0, :] *= render_width / width Ks[..., 1, :] *= render_height / height - torch.cuda.reset_peak_memory_stats() - mem_tic = torch.cuda.max_memory_allocated() / 1024**3 + torch_acc.reset_peak_memory_stats() + mem_tic = torch_acc.max_memory_allocated() / 1024**3 if memory_history: torch.cuda.memory._record_memory_history() @@ -105,7 +106,7 @@ def main( with_ut=model == "3DGUT", with_eval3d=model == "3DGUT", ) - mem_toc_fwd = torch.cuda.max_memory_allocated() / 1024**3 - mem_tic + mem_toc_fwd = torch_acc.max_memory_allocated() / 1024**3 - mem_tic render_colors = outputs[0] loss = render_colors.sum() @@ -116,7 +117,7 @@ def backward(): v.grad = None ellipse_time_bwd, _ = timeit(repeats, backward) - mem_toc_all = torch.cuda.max_memory_allocated() / 1024**3 - mem_tic + mem_toc_all = torch_acc.max_memory_allocated() / 1024**3 - mem_tic print( f"Rasterization Mem Allocation: [FWD]{mem_toc_fwd:.2f} GB, [All]{mem_toc_all:.2f} GB " f"Time: [FWD]{ellipse_time_fwd:.3f}s, [BWD]{ellipse_time_bwd:.3f}s " @@ -176,7 +177,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): f"{stats['time_bwd']:0.5f}", ] ) - torch.cuda.empty_cache() + torch_acc.empty_cache() if world_rank == 0: headers = [ @@ -269,5 +270,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): args = parser.parse_args() if args.memory_history: args.repeats = 1 # only run once for memory history + if BACKEND != "cuda": + raise ValueError("Memory history is only supported for CUDA backend.") cli(worker, args, verbose=True) diff --git a/profiling/main.py b/profiling/main.py index e8d7c24e..e4e636c1 100644 --- a/profiling/main.py +++ b/profiling/main.py @@ -6,11 +6,13 @@ ``` """ +import os import time import torch from typing_extensions import Callable, Literal +from gsplat import __version__, torch_acc, BACKEND from gsplat._helper import load_test_data from gsplat.distributed import cli from gsplat.rendering import rasterization @@ -22,17 +24,17 @@ "4k": (3840, 2160), } -device = torch.device("cuda") +device = torch_acc._device(0) def timeit(repeats: int, f: Callable, *args, **kwargs) -> float: for _ in range(5): # warmup f(*args, **kwargs) - torch.cuda.synchronize() + torch_acc.synchronize() start = time.time() for _ in range(repeats): results = f(*args, **kwargs) - torch.cuda.synchronize() + torch_acc.synchronize() end = time.time() return (end - start) / repeats, results @@ -50,6 +52,7 @@ def main( world_rank: int = 0, world_size: int = 1, ): + data_path = os.path.join(os.path.dirname(__file__), "../assets/test_garden.npz") ( means, quats, @@ -60,8 +63,7 @@ def main( Ks, width, height, - ) = load_test_data(device=device, scene_grid=scene_grid) - + ) = load_test_data(data_path=data_path, device=device, scene_grid=scene_grid) # to batch viewmats = viewmats[:1].repeat(batch_size, 1, 1) Ks = Ks[:1].repeat(batch_size, 1, 1) @@ -86,8 +88,8 @@ def main( Ks[..., 0, :] *= render_width / width Ks[..., 1, :] *= render_height / height - torch.cuda.reset_peak_memory_stats() - mem_tic = torch.cuda.max_memory_allocated() / 1024**3 + torch_acc.reset_peak_memory_stats() + mem_tic = torch_acc.max_memory_allocated() / 1024**3 if memory_history: torch.cuda.memory._record_memory_history() @@ -120,7 +122,7 @@ def main( sparse_grad=sparse_grad, distributed=world_size > 1, ) - mem_toc_fwd = torch.cuda.max_memory_allocated() / 1024**3 - mem_tic + mem_toc_fwd = torch_acc.max_memory_allocated() / 1024**3 - mem_tic render_colors = outputs[0] loss = render_colors.sum() @@ -131,7 +133,7 @@ def backward(): v.grad = None ellipse_time_bwd, _ = timeit(repeats, backward) - mem_toc_all = torch.cuda.max_memory_allocated() / 1024**3 - mem_tic + mem_toc_all = torch_acc.max_memory_allocated() / 1024**3 - mem_tic print( f"Rasterization Mem Allocation: [FWD]{mem_toc_fwd:.2f} GB, [All]{mem_toc_all:.2f} GB " f"Time: [FWD]{ellipse_time_fwd:.3f}s, [BWD]{ellipse_time_bwd:.3f}s " @@ -180,7 +182,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): ) collection.append( [ - "gsplat v1.0.0", + f"gsplat v{__version__}", True, True, # configs @@ -194,7 +196,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): f"{1.0 / stats['time_bwd']:0.1f} x {(batch_size)}", ] ) - torch.cuda.empty_cache() + torch_acc.empty_cache() print("gsplat packed[True] sparse_grad[False]") for scene_grid in args.scene_grid: @@ -211,7 +213,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): ) collection.append( [ - "gsplat v1.0.0", + f"gsplat v{__version__}", True, False, # configs @@ -225,7 +227,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): f"{1.0 / stats['time_bwd']:0.1f} x {(batch_size)}", ] ) - torch.cuda.empty_cache() + torch_acc.empty_cache() print("gsplat packed[False] sparse_grad[False]") for scene_grid in args.scene_grid: @@ -242,7 +244,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): ) collection.append( [ - "gsplat v1.0.0", + f"gsplat v{__version__}", False, False, # configs @@ -256,7 +258,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): f"{1.0 / stats['time_bwd']:0.1f} x {(batch_size)}", ] ) - torch.cuda.empty_cache() + torch_acc.empty_cache() if "inria" in args.backends: print("inria") @@ -285,7 +287,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): f"{1.0 / stats['time_bwd']:0.1f} x {(batch_size)}", ] ) - torch.cuda.empty_cache() + torch_acc.empty_cache() if world_rank == 0: headers = [ @@ -366,5 +368,7 @@ def worker(local_rank: int, world_rank: int, world_size: int, args): args = parser.parse_args() if args.memory_history: args.repeats = 1 # only run once for memory history + if BACKEND != "cuda": + raise ValueError("Memory history is only supported for CUDA backend.") cli(worker, args, verbose=True) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 00000000..61b6882b --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "torch", "ninja", "cmake", "pybind11>=2.10"] +build-backend = "setuptools.build_meta" \ No newline at end of file diff --git a/setup.py b/setup.py index c0ba3c48..0a9a78e4 100644 --- a/setup.py +++ b/setup.py @@ -12,7 +12,16 @@ URL = "https://github.com/nerfstudio-project/gsplat" -BUILD_NO_CUDA = os.getenv("BUILD_NO_CUDA", "0") == "1" +has_xpu = False +try: + import torch + + has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() +except (ImportError, AttributeError): + pass + +BUILD_SYCL = has_xpu or os.getenv("BUILD_SYCL", "0") == "1" +BUILD_NO_CUDA = os.getenv("BUILD_NO_CUDA", "0") == "1" or BUILD_SYCL WITH_SYMBOLS = os.getenv("WITH_SYMBOLS", "0") == "1" LINE_INFO = os.getenv("LINE_INFO", "0") == "1" MAX_JOBS = os.getenv("MAX_JOBS") @@ -26,7 +35,49 @@ def get_ext(): from torch.utils.cpp_extension import BuildExtension - return BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=True) + if not BUILD_NO_CUDA: + return BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=True) + if not BUILD_SYCL: + return None + + class SyclBuildExtension(BuildExtension): + """ + Custom build class to orchestrate a CMake build for the SYCL backend. + """ + + def run(self): + print("--- Running SYCL build via CMake ---") + import shutil + import subprocess as sp + + sycl_dir = os.path.abspath("gsplat/sycl") + build_dir = os.path.join(self.build_temp, "sycl") + os.makedirs(build_dir, exist_ok=True) + jobs = os.getenv("MAX_JOBS", "10") + install_dir = os.path.abspath(self.build_lib) + + cfg = "RelWithDebInfo" if WITH_SYMBOLS or LINE_INFO else "Release" + cmake_args = [ + "cmake", + "-G", + "Ninja", + f"-DCMAKE_BUILD_TYPE={cfg}", + f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={os.path.join(install_dir, 'gsplat')}", + sycl_dir, + ] + # ninja not found in isolated env during "python -m build" + if ninja_path := shutil.which("ninja"): + cmake_args.append(f"-DCMAKE_MAKE_PROGRAM={ninja_path}") + sp.check_call( + cmake_args, + cwd=build_dir, + ) + sp.check_call( + ["cmake", "--build", ".", "--config", cfg, "--", "-v", f"-j{jobs}"], + cwd=build_dir, + ) + + return SyclBuildExtension def get_extensions(): @@ -105,14 +156,30 @@ def get_extensions(): return [extension] +ext_modules = [] +cmdclass = {} +packages_to_find = find_packages() +from setuptools import Extension + +if BUILD_SYCL: + print("--- Configuring for SYCL build ---") + cmdclass = {"build_ext": get_ext()} + ext_modules.append(Extension("gsplat.gsplat_sycl_kernels", sources=[])) +elif not BUILD_NO_CUDA: + print("--- Configuring for CUDA build ---") + cmdclass = {"build_ext": get_ext()} + ext_modules = get_extensions() +else: + print("--- Building without any C++/CUDA/SYCL extensions ---") + setup( name="gsplat", version=__version__, - description=" Python package for differentiable rasterization of gaussians", - keywords="gaussian, splatting, cuda", + description="Python package for differentiable rasterization of gaussians", + keywords="gaussian, splatting, cuda, sycl", url=URL, download_url=f"{URL}/archive/gsplat-{__version__}.tar.gz", - python_requires=">=3.7", + python_requires=">=3.8", install_requires=[ "ninja", "numpy", @@ -135,11 +202,11 @@ def get_extensions(): "twine", ], }, - ext_modules=get_extensions() if not BUILD_NO_CUDA else [], - cmdclass={"build_ext": get_ext()} if not BUILD_NO_CUDA else {}, - packages=find_packages(), - # https://github.com/pypa/setuptools/issues/1461#issuecomment-954725244 + ext_modules=ext_modules, + cmdclass=cmdclass, + packages=packages_to_find, include_package_data=True, + zip_safe=False, ) if need_to_unset_max_jobs: diff --git a/tests/_test_distributed.py b/tests/_test_distributed.py index ae03f9c9..46de2b89 100644 --- a/tests/_test_distributed.py +++ b/tests/_test_distributed.py @@ -1,6 +1,7 @@ import pytest import torch +import gsplat from gsplat.distributed import ( all_gather_int32, all_gather_tensor_list, @@ -9,9 +10,14 @@ cli, ) +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), + reason="No CUDA or SYCL XPU backend available", +) + def _main_all_gather_int32(local_rank: int, world_rank: int, world_size: int, _): - device = torch.device("cuda", local_rank) + device = torch.device(local_rank) value = world_rank collected = all_gather_int32(world_size, value, device=device) @@ -24,13 +30,13 @@ def _main_all_gather_int32(local_rank: int, world_rank: int, world_size: int, _) assert collected[i] == torch.tensor(i, device=device, dtype=torch.int) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend def test_all_gather_int32(): cli(_main_all_gather_int32, None, verbose=True) def _main_all_to_all_int32(local_rank: int, world_rank: int, world_size: int, _): - device = torch.device("cuda", local_rank) + device = torch.device(local_rank) values = list(range(world_size)) collected = all_to_all_int32(world_size, values, device=device) @@ -43,13 +49,13 @@ def _main_all_to_all_int32(local_rank: int, world_rank: int, world_size: int, _) assert collected[i] == torch.tensor(world_rank, device=device, dtype=torch.int) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend def test_all_to_all_int32(): cli(_main_all_to_all_int32, None, verbose=True) def _main_all_gather_tensor_list(local_rank: int, world_rank: int, world_size: int, _): - device = torch.device("cuda", local_rank) + device = torch.device(local_rank) N = 10 tensor_list = [ @@ -67,13 +73,13 @@ def _main_all_gather_tensor_list(local_rank: int, world_rank: int, world_size: i assert torch.equal(tensor, target) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend def test_all_gather_tensor_list(): cli(_main_all_gather_tensor_list, None, verbose=True) def _main_all_to_all_tensor_list(local_rank: int, world_rank: int, world_size: int, _): - device = torch.device("cuda", local_rank) + device = torch.device(local_rank) splits = torch.arange(0, world_size, device=device) N = splits.sum().item() @@ -102,7 +108,7 @@ def _main_all_to_all_tensor_list(local_rank: int, world_rank: int, world_size: i assert torch.equal(tensor, target) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend def test_all_to_all_tensor_list(): cli(_main_all_to_all_tensor_list, None, verbose=True) diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 788c9111..9475b617 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -4,7 +4,21 @@ import torch from typing_extensions import Tuple -device = torch.device("cuda:0") +import gsplat + +if gsplat.BACKEND == "sycl": + device = torch.device("xpu:0") +elif gsplat.BACKEND == "cuda": + device = torch.device("cuda:0") +else: + device = torch.device("cpu") + +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" +) +requires_cuda = pytest.mark.skipif( + gsplat.BACKEND != "cuda", reason="No CUDA backend available" +) def expand(data: dict, batch_dims: Tuple[int, ...]): @@ -22,7 +36,6 @@ def expand(data: dict, batch_dims: Tuple[int, ...]): @pytest.fixture -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") def test_data(): C = 3 N = 1000 @@ -54,11 +67,11 @@ def test_data(): } -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_projection_2dgs(test_data, batch_dims: Tuple[int, ...]): - from gsplat.cuda._torch_impl_2dgs import _fully_fused_projection_2dgs - from gsplat.cuda._wrapper import fully_fused_projection_2dgs + from gsplat._torch_impl_2dgs import _fully_fused_projection_2dgs + from gsplat import fully_fused_projection_2dgs torch.manual_seed(42) @@ -124,13 +137,13 @@ def test_projection_2dgs(test_data, batch_dims: Tuple[int, ...]): torch.testing.assert_close(v_means, _v_means, rtol=1e-2, atol=6e-2) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_cuda @pytest.mark.parametrize("sparse_grad", [False]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_fully_fused_projection_packed_2dgs( test_data, sparse_grad: bool, batch_dims: Tuple[int, ...] ): - from gsplat.cuda._wrapper import fully_fused_projection_2dgs + from gsplat._wrapper import fully_fused_projection_2dgs torch.manual_seed(42) @@ -248,14 +261,14 @@ def test_fully_fused_projection_packed_2dgs( torch.testing.assert_close(v_quats, _v_quats, rtol=1e-2, atol=1e-2) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("channels", [3, 31]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_rasterize_to_pixels_2dgs( test_data, channels: int, batch_dims: Tuple[int, ...] ): - from gsplat.cuda._torch_impl_2dgs import _rasterize_to_pixels_2dgs - from gsplat.cuda._wrapper import ( + from gsplat._torch_impl_2dgs import _rasterize_to_pixels_2dgs + from gsplat import ( fully_fused_projection_2dgs, isect_offset_encode, isect_tiles, @@ -327,19 +340,20 @@ def test_rasterize_to_pixels_2dgs( distloss=True, ) - _render_colors, _render_alphas, _render_normals = _rasterize_to_pixels_2dgs( - means2d, - ray_transforms, - colors, - normals, - opacities, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - backgrounds=backgrounds, - ) + if gsplat.BACKEND != "sycl": # nerfacc required for comparison + _render_colors, _render_alphas, _render_normals = _rasterize_to_pixels_2dgs( + means2d, + ray_transforms, + colors, + normals, + opacities, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + backgrounds=backgrounds, + ) v_render_colors = torch.rand_like(render_colors) v_render_alphas = torch.rand_like(render_alphas) @@ -359,31 +373,34 @@ def test_rasterize_to_pixels_2dgs( (means2d, ray_transforms, colors, opacities, backgrounds, normals), ) - ( - _v_means2d, - _v_ray_transforms, - _v_colors, - _v_opacities, - _v_backgrounds, - _v_normals, - ) = torch.autograd.grad( - (_render_colors * v_render_colors).sum() - + (_render_alphas * v_render_alphas).sum() - + (_render_normals * v_render_normals).sum(), - (means2d, ray_transforms, colors, opacities, backgrounds, normals), - ) - - # assert close forward - torch.testing.assert_close(render_colors, _render_colors, atol=1e-3, rtol=1e-3) - torch.testing.assert_close(render_alphas, _render_alphas, atol=1e-3, rtol=1e-3) - torch.testing.assert_close(render_normals, _render_normals, atol=1e-3, rtol=1e-3) - - # assert close backward - torch.testing.assert_close(v_means2d, _v_means2d, rtol=1e-3, atol=1e-3) - torch.testing.assert_close( - v_ray_transforms, _v_ray_transforms, rtol=2e-1, atol=5e-2 - ) - torch.testing.assert_close(v_colors, _v_colors, rtol=1e-3, atol=1e-3) - torch.testing.assert_close(v_opacities, _v_opacities, rtol=1e-3, atol=1e-3) - torch.testing.assert_close(v_backgrounds, _v_backgrounds, rtol=1e-5, atol=1e-5) - torch.testing.assert_close(v_normals, _v_normals, rtol=1e-3, atol=1e-3) + if gsplat.BACKEND != "sycl": # nerfacc required for comparison + ( + _v_means2d, + _v_ray_transforms, + _v_colors, + _v_opacities, + _v_backgrounds, + _v_normals, + ) = torch.autograd.grad( + (_render_colors * v_render_colors).sum() + + (_render_alphas * v_render_alphas).sum() + + (_render_normals * v_render_normals).sum(), + (means2d, ray_transforms, colors, opacities, backgrounds, normals), + ) + + # assert close forward + torch.testing.assert_close(render_colors, _render_colors, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(render_alphas, _render_alphas, atol=1e-3, rtol=1e-3) + torch.testing.assert_close( + render_normals, _render_normals, atol=1e-3, rtol=1e-3 + ) + + # assert close backward + torch.testing.assert_close(v_means2d, _v_means2d, rtol=1e-3, atol=1e-3) + torch.testing.assert_close( + v_ray_transforms, _v_ray_transforms, rtol=2e-1, atol=5e-2 + ) + torch.testing.assert_close(v_colors, _v_colors, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(v_opacities, _v_opacities, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(v_backgrounds, _v_backgrounds, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(v_normals, _v_normals, rtol=1e-3, atol=1e-3) diff --git a/tests/test_basic.py b/tests/test_basic.py index 1849f288..0cba2926 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,7 +1,7 @@ -"""Tests for the functions in the CUDA extension. +"""Tests for the backend functions. Usage: -```bash + pytest -s ``` """ @@ -13,15 +13,28 @@ import torch from typing_extensions import Literal, Tuple, assert_never +# Import the gsplat library, which will run the __init__.py and select a backend. +import gsplat from gsplat._helper import load_test_data -device = torch.device("cuda:0") +if gsplat.BACKEND == "sycl": + device = torch.device("xpu:0") +elif gsplat.BACKEND == "cuda": + device = torch.device("cuda:0") +else: + device = None + +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), + reason="No CUDA or SYCL XPU backend available", +) +requires_cuda = pytest.mark.skipif( + gsplat.BACKEND != "cuda", reason="Test requires CUDA backend" +) def expand(data: dict, batch_dims: Tuple[int, ...]): - # append multiple batch dimensions to the front of the tensor - # eg. x.shape = [N, 3], batch_dims = (1, 2), return shape is [1, 2, N, 3] - # eg. x.shape = [N, 3], batch_dims = (), return shape is [N, 3] + """Helper function to expand test data with batch dimensions.""" ret = {} for k, v in data.items(): if isinstance(v, torch.Tensor) and len(batch_dims) > 0: @@ -34,6 +47,7 @@ def expand(data: dict, batch_dims: Tuple[int, ...]): @pytest.fixture def test_data(): + """Loads test data and moves it to the active device.""" ( means, quats, @@ -49,89 +63,76 @@ def test_data(): data_path=os.path.join(os.path.dirname(__file__), "../assets/test_garden.npz"), ) return { - "means": means, # [N, 3] - "quats": quats, # [N, 4] - "scales": scales, # [N, 3] - "opacities": opacities, # [N] - "viewmats": viewmats, # [C, 4, 4] - "Ks": Ks, # [C, 3, 3] + "means": means, + "quats": quats, + "scales": scales, + "opacities": opacities, + "viewmats": viewmats, + "Ks": Ks, "width": width, "height": height, } -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("triu", [False, True]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_quat_scale_to_covar_preci(test_data, triu: bool, batch_dims: Tuple[int, ...]): - from gsplat.cuda._torch_impl import _quat_scale_to_covar_preci - from gsplat.cuda._wrapper import quat_scale_to_covar_preci - torch.manual_seed(42) + from gsplat._torch_impl import _quat_scale_to_covar_preci + torch.manual_seed(42) test_data = expand(test_data, batch_dims) quats = test_data["quats"] scales = test_data["scales"] quats.requires_grad = True scales.requires_grad = True - # forward - covars, precis = quat_scale_to_covar_preci(quats, scales, triu=triu) + covars, precis = gsplat.quat_scale_to_covar_preci(quats, scales, triu=triu) _covars, _precis = _quat_scale_to_covar_preci(quats, scales, triu=triu) torch.testing.assert_close(covars, _covars) - # This test is disabled because the numerical instability. - # torch.testing.assert_close(precis, _precis, rtol=2e-2, atol=1e-2) - # if not triu: - # I = torch.eye(3, device=device).expand(len(covars), 3, 3) - # torch.testing.assert_close(torch.bmm(covars, precis), I) - # torch.testing.assert_close(torch.bmm(precis, covars), I) - - # backward + v_covars = torch.randn_like(covars) v_precis = torch.randn_like(precis) * 0.01 v_quats, v_scales = torch.autograd.grad( - (covars * v_covars + precis * v_precis).sum(), - (quats, scales), + (covars * v_covars + precis * v_precis).sum(), (quats, scales) ) _v_quats, _v_scales = torch.autograd.grad( - (_covars * v_covars + _precis * v_precis).sum(), - (quats, scales), + (_covars * v_covars + _precis * v_precis).sum(), (quats, scales) ) torch.testing.assert_close(v_quats, _v_quats, rtol=1e0, atol=1e-1) torch.testing.assert_close(v_scales, _v_scales, rtol=1e0, atol=1e-1) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("camera_model", ["pinhole", "ortho", "fisheye"]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) -def test_proj( - test_data, - camera_model: Literal["pinhole", "ortho", "fisheye"], - batch_dims: Tuple[int, ...], -): - from gsplat.cuda._torch_impl import ( +def test_proj(test_data, camera_model: str, batch_dims: Tuple[int, ...]): + + from gsplat._torch_impl import ( _fisheye_proj, _ortho_proj, _persp_proj, _world_to_cam, ) - from gsplat.cuda._wrapper import proj, quat_scale_to_covar_preci torch.manual_seed(42) - test_data = expand(test_data, batch_dims) - Ks = test_data["Ks"] - viewmats = test_data["viewmats"] - height = test_data["height"] - width = test_data["width"] + Ks, viewmats, height, width = ( + test_data["Ks"], + test_data["viewmats"], + test_data["height"], + test_data["width"], + ) - covars, _ = quat_scale_to_covar_preci(test_data["quats"], test_data["scales"]) + covars, _ = gsplat.quat_scale_to_covar_preci( + test_data["quats"], test_data["scales"] + ) means, covars = _world_to_cam(test_data["means"], covars, viewmats) means.requires_grad = True covars.requires_grad = True - # forward - means2d, covars2d = proj(means, covars, Ks, width, height, camera_model) + means2d, covars2d = gsplat.proj(means, covars, Ks, width, height, camera_model) if camera_model == "ortho": _means2d, _covars2d = _ortho_proj(means, covars, Ks, width, height) elif camera_model == "fisheye": @@ -144,22 +145,18 @@ def test_proj( torch.testing.assert_close(means2d, _means2d, rtol=1e-4, atol=1e-4) torch.testing.assert_close(covars2d, _covars2d, rtol=1e-1, atol=3e-2) - # backward - v_means2d = torch.randn_like(means2d) - v_covars2d = torch.randn_like(covars2d) + v_means2d, v_covars2d = torch.randn_like(means2d), torch.randn_like(covars2d) v_means, v_covars = torch.autograd.grad( - (means2d * v_means2d).sum() + (covars2d * v_covars2d).sum(), - (means, covars), + (means2d * v_means2d).sum() + (covars2d * v_covars2d).sum(), (means, covars) ) _v_means, _v_covars = torch.autograd.grad( - (_means2d * v_means2d).sum() + (_covars2d * v_covars2d).sum(), - (means, covars), + (_means2d * v_means2d).sum() + (_covars2d * v_covars2d).sum(), (means, covars) ) torch.testing.assert_close(v_means, _v_means, rtol=6e-1, atol=1e-2) torch.testing.assert_close(v_covars, _v_covars, rtol=1e-1, atol=1e-1) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("camera_model", ["pinhole", "ortho", "fisheye"]) @pytest.mark.parametrize("fused", [False, True]) @pytest.mark.parametrize("calc_compensations", [True, False]) @@ -168,31 +165,28 @@ def test_projection( test_data, fused: bool, calc_compensations: bool, - camera_model: Literal["pinhole", "ortho", "fisheye"], + camera_model: str, batch_dims: Tuple[int, ...], ): - from gsplat.cuda._torch_impl import _fully_fused_projection - from gsplat.cuda._wrapper import fully_fused_projection, quat_scale_to_covar_preci - torch.manual_seed(42) + from gsplat._torch_impl import _fully_fused_projection + torch.manual_seed(42) test_data = expand(test_data, batch_dims) - Ks = test_data["Ks"] - viewmats = test_data["viewmats"] - height = test_data["height"] - width = test_data["width"] - quats = test_data["quats"] - scales = test_data["scales"] - means = test_data["means"] - + Ks, viewmats, height, width = ( + test_data["Ks"], + test_data["viewmats"], + test_data["height"], + test_data["width"], + ) + quats, scales, means = test_data["quats"], test_data["scales"], test_data["means"] viewmats.requires_grad = True quats.requires_grad = True scales.requires_grad = True means.requires_grad = True - # forward if fused: - radii, means2d, depths, conics, compensations = fully_fused_projection( + radii, means2d, depths, conics, compensations = gsplat.fully_fused_projection( means, None, quats, @@ -205,8 +199,8 @@ def test_projection( camera_model=camera_model, ) else: - covars, _ = quat_scale_to_covar_preci(quats, scales, triu=True) # [..., N, 6] - radii, means2d, depths, conics, compensations = fully_fused_projection( + covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=True) + radii, means2d, depths, conics, compensations = gsplat.fully_fused_projection( means, covars, None, @@ -218,7 +212,8 @@ def test_projection( calc_compensations=calc_compensations, camera_model=camera_model, ) - _covars, _ = quat_scale_to_covar_preci(quats, scales, triu=False) # [..., N, 3, 3] + + _covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=False) _radii, _means2d, _depths, _conics, _compensations = _fully_fused_projection( means, _covars, @@ -230,7 +225,6 @@ def test_projection( camera_model=camera_model, ) - # radii is integer so we allow for 1 unit difference valid = (radii > 0).all(dim=-1) & (_radii > 0).all(dim=-1) torch.testing.assert_close(radii, _radii, rtol=0, atol=1) torch.testing.assert_close(means2d[valid], _means2d[valid], rtol=1e-4, atol=1e-4) @@ -241,25 +235,32 @@ def test_projection( compensations[valid], _compensations[valid], rtol=1e-4, atol=1e-3 ) - # backward - v_means2d = torch.randn_like(means2d) * valid[..., None] - v_depths = torch.randn_like(depths) * valid - v_conics = torch.randn_like(conics) * valid[..., None] - if calc_compensations: - v_compensations = torch.randn_like(compensations) * valid - v_viewmats, v_quats, v_scales, v_means = torch.autograd.grad( + v_means2d, v_depths, v_conics = ( + torch.randn_like(means2d) * valid[..., None], + torch.randn_like(depths) * valid, + torch.randn_like(conics) * valid[..., None], + ) + v_compensations = ( + torch.randn_like(compensations) * valid if calc_compensations else 0 + ) + grad_sum = ( (means2d * v_means2d).sum() + (depths * v_depths).sum() + (conics * v_conics).sum() - + ((compensations * v_compensations).sum() if calc_compensations else 0), - (viewmats, quats, scales, means), + + ((compensations * v_compensations).sum() if calc_compensations else 0) ) - _v_viewmats, _v_quats, _v_scales, _v_means = torch.autograd.grad( + v_viewmats, v_quats, v_scales, v_means = torch.autograd.grad( + grad_sum, (viewmats, quats, scales, means) + ) + + _grad_sum = ( (_means2d * v_means2d).sum() + (_depths * v_depths).sum() + (_conics * v_conics).sum() - + ((_compensations * v_compensations).sum() if calc_compensations else 0), - (viewmats, quats, scales, means), + + ((_compensations * v_compensations).sum() if calc_compensations else 0) + ) + _v_viewmats, _v_quats, _v_scales, _v_means = torch.autograd.grad( + _grad_sum, (viewmats, quats, scales, means) ) torch.testing.assert_close(v_viewmats, _v_viewmats, rtol=2e-3, atol=2e-3) @@ -268,7 +269,7 @@ def test_projection( torch.testing.assert_close(v_means, _v_means, rtol=1e-2, atol=6e-2) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("fused", [False, True]) @pytest.mark.parametrize("sparse_grad", [False]) @pytest.mark.parametrize("calc_compensations", [False, True]) @@ -279,28 +280,24 @@ def test_fully_fused_projection_packed( fused: bool, sparse_grad: bool, calc_compensations: bool, - camera_model: Literal["pinhole", "ortho", "fisheye"], + camera_model: str, batch_dims: Tuple[int, ...], ): - from gsplat.cuda._wrapper import fully_fused_projection, quat_scale_to_covar_preci torch.manual_seed(42) - test_data = expand(test_data, batch_dims) - Ks = test_data["Ks"] - viewmats = test_data["viewmats"] - height = test_data["height"] - width = test_data["width"] - quats = test_data["quats"] - scales = test_data["scales"] - means = test_data["means"] - + Ks, viewmats, height, width = ( + test_data["Ks"], + test_data["viewmats"], + test_data["height"], + test_data["width"], + ) + quats, scales, means = test_data["quats"], test_data["scales"], test_data["means"] viewmats.requires_grad = True quats.requires_grad = True scales.requires_grad = True means.requires_grad = True - # forward if fused: ( batch_ids, @@ -312,7 +309,7 @@ def test_fully_fused_projection_packed( depths, conics, compensations, - ) = fully_fused_projection( + ) = gsplat.fully_fused_projection( means, None, quats, @@ -326,7 +323,13 @@ def test_fully_fused_projection_packed( calc_compensations=calc_compensations, camera_model=camera_model, ) - _radii, _means2d, _depths, _conics, _compensations = fully_fused_projection( + ( + _radii, + _means2d, + _depths, + _conics, + _compensations, + ) = gsplat.fully_fused_projection( means, None, quats, @@ -340,7 +343,9 @@ def test_fully_fused_projection_packed( camera_model=camera_model, ) else: - covars, _ = quat_scale_to_covar_preci(quats, scales, triu=True) # [..., N, 6] + covars, _ = gsplat.quat_scale_to_covar_preci( + quats, scales, triu=True + ) # [..., N, 6] ( batch_ids, camera_ids, @@ -351,7 +356,7 @@ def test_fully_fused_projection_packed( depths, conics, compensations, - ) = fully_fused_projection( + ) = gsplat.fully_fused_projection( means, covars, None, @@ -365,7 +370,13 @@ def test_fully_fused_projection_packed( calc_compensations=calc_compensations, camera_model=camera_model, ) - _radii, _means2d, _depths, _conics, _compensations = fully_fused_projection( + ( + _radii, + _means2d, + _depths, + _conics, + _compensations, + ) = gsplat.fully_fused_projection( means, covars, None, @@ -379,34 +390,38 @@ def test_fully_fused_projection_packed( camera_model=camera_model, ) - B = math.prod(batch_dims) - N = means.shape[-2] - C = viewmats.shape[-3] - - # recover packed tensors to full matrices for testing - __radii = torch.sparse_coo_tensor( - torch.stack([batch_ids, camera_ids, gaussian_ids]), radii, (B, C, N, 2) - ).to_dense() - __radii = __radii.reshape(batch_dims + (C, N, 2)) - __means2d = torch.sparse_coo_tensor( - torch.stack([batch_ids, camera_ids, gaussian_ids]), means2d, (B, C, N, 2) - ).to_dense() - __means2d = __means2d.reshape(batch_dims + (C, N, 2)) - __depths = torch.sparse_coo_tensor( - torch.stack([batch_ids, camera_ids, gaussian_ids]), depths, (B, C, N) - ).to_dense() - __depths = __depths.reshape(batch_dims + (C, N)) - __conics = torch.sparse_coo_tensor( - torch.stack([batch_ids, camera_ids, gaussian_ids]), conics, (B, C, N, 3) - ).to_dense() - __conics = __conics.reshape(batch_dims + (C, N, 3)) + B, C, N = math.prod(batch_dims), viewmats.shape[-3], means.shape[-2] + + # Unpack for comparison + sparse_shape = (B, C, N) + indices = torch.stack([batch_ids, camera_ids, gaussian_ids]) + __radii = ( + torch.sparse_coo_tensor(indices, radii, sparse_shape + (2,)) + .to_dense() + .reshape(batch_dims + (C, N, 2)) + ) + __means2d = ( + torch.sparse_coo_tensor(indices, means2d, sparse_shape + (2,)) + .to_dense() + .reshape(batch_dims + (C, N, 2)) + ) + __depths = ( + torch.sparse_coo_tensor(indices, depths, sparse_shape) + .to_dense() + .reshape(batch_dims + (C, N)) + ) + __conics = ( + torch.sparse_coo_tensor(indices, conics, sparse_shape + (3,)) + .to_dense() + .reshape(batch_dims + (C, N, 3)) + ) if calc_compensations: - __compensations = torch.sparse_coo_tensor( - torch.stack([batch_ids, camera_ids, gaussian_ids]), - compensations, - (B, C, N), - ).to_dense() - __compensations = __compensations.reshape(batch_dims + (C, N)) + __compensations = ( + torch.sparse_coo_tensor(indices, compensations, sparse_shape) + .to_dense() + .reshape(batch_dims + (C, N)) + ) + sel = (__radii > 0).all(dim=-1) & (_radii > 0).all(dim=-1) torch.testing.assert_close(__radii[sel], _radii[sel], rtol=0, atol=1) torch.testing.assert_close(__means2d[sel], _means2d[sel], rtol=1e-4, atol=1e-4) @@ -417,10 +432,11 @@ def test_fully_fused_projection_packed( __compensations[sel], _compensations[sel], rtol=1e-4, atol=1e-3 ) - # backward - v_means2d = torch.randn_like(_means2d) * sel[..., None] - v_depths = torch.randn_like(_depths) * sel - v_conics = torch.randn_like(_conics) * sel[..., None] + v_means2d, v_depths, v_conics = ( + torch.randn_like(_means2d) * sel[..., None], + torch.randn_like(_depths) * sel, + torch.randn_like(_conics) * sel[..., None], + ) _v_viewmats, _v_quats, _v_scales, _v_means = torch.autograd.grad( (_means2d * v_means2d).sum() + (_depths * v_depths).sum() @@ -429,16 +445,18 @@ def test_fully_fused_projection_packed( retain_graph=True, ) v_viewmats, v_quats, v_scales, v_means = torch.autograd.grad( - (means2d * v_means2d[(__radii > 0).all(dim=-1)]).sum() - + (depths * v_depths[(__radii > 0).all(dim=-1)]).sum() - + (conics * v_conics[(__radii > 0).all(dim=-1)]).sum(), + (means2d * v_means2d[sel]).sum() + + (depths * v_depths[sel]).sum() + + (conics * v_conics[sel]).sum(), (viewmats, quats, scales, means), retain_graph=True, ) if sparse_grad: - v_quats = v_quats.to_dense() - v_scales = v_scales.to_dense() - v_means = v_means.to_dense() + v_quats, v_scales, v_means = ( + v_quats.to_dense(), + v_scales.to_dense(), + v_means.to_dense(), + ) torch.testing.assert_close(v_viewmats, _v_viewmats, rtol=1e-2, atol=1e-2) torch.testing.assert_close(v_quats, _v_quats, rtol=1e-3, atol=1e-3) @@ -446,18 +464,15 @@ def test_fully_fused_projection_packed( torch.testing.assert_close(v_means, _v_means, rtol=1e-3, atol=1e-3) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_isect(test_data, batch_dims: Tuple[int, ...]): - from gsplat.cuda._torch_impl import _isect_offset_encode, _isect_tiles - from gsplat.cuda._wrapper import isect_offset_encode, isect_tiles - torch.manual_seed(42) + from gsplat._torch_impl import _isect_offset_encode, _isect_tiles - B = math.prod(batch_dims) - C, N = 3, 1000 - I = B * C - width, height = 40, 60 + torch.manual_seed(42) + B, C, N = math.prod(batch_dims), 3, 1000 + I, width, height = B * C, 40, 60 test_data = { "means2d": torch.randn(C, N, 2, device=device) * width, @@ -465,18 +480,21 @@ def test_isect(test_data, batch_dims: Tuple[int, ...]): "depths": torch.rand(C, N, device=device), } test_data = expand(test_data, batch_dims) - means2d = test_data["means2d"] - radii = test_data["radii"] - depths = test_data["depths"] + means2d, radii, depths = ( + test_data["means2d"], + test_data["radii"], + test_data["depths"], + ) tile_size = 16 - tile_width = math.ceil(width / tile_size) - tile_height = math.ceil(height / tile_size) + tile_width, tile_height = math.ceil(width / tile_size), math.ceil( + height / tile_size + ) - tiles_per_gauss, isect_ids, flatten_ids = isect_tiles( + tiles_per_gauss, isect_ids, flatten_ids = gsplat.isect_tiles( means2d, radii, depths, tile_size, tile_width, tile_height ) - isect_offsets = isect_offset_encode(isect_ids, I, tile_width, tile_height) + isect_offsets = gsplat.isect_offset_encode(isect_ids, I, tile_width, tile_height) _tiles_per_gauss, _isect_ids, _gauss_ids = _isect_tiles( means2d, radii, depths, tile_size, tile_width, tile_height @@ -489,23 +507,15 @@ def test_isect(test_data, batch_dims: Tuple[int, ...]): torch.testing.assert_close(isect_offsets, _isect_offsets) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend @pytest.mark.parametrize("channels", [3, 32, 128]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, ...]): - from gsplat.cuda._torch_impl import _rasterize_to_pixels - from gsplat.cuda._wrapper import ( - fully_fused_projection, - isect_offset_encode, - isect_tiles, - quat_scale_to_covar_preci, - rasterize_to_pixels, - ) - torch.manual_seed(42) + from gsplat._torch_impl import _rasterize_to_pixels - N = test_data["means"].shape[-2] - C = test_data["viewmats"].shape[-3] + torch.manual_seed(42) + N, C = test_data["means"].shape[-2], test_data["viewmats"].shape[-3] I = math.prod(batch_dims) * C test_data.update( { @@ -514,34 +524,38 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. } ) test_data = expand(test_data, batch_dims) - Ks = test_data["Ks"] - viewmats = test_data["viewmats"] - height = test_data["height"] - width = test_data["width"] - quats = test_data["quats"] - scales = test_data["scales"] * 0.1 - means = test_data["means"] - opacities = test_data["opacities"] - colors = test_data["colors"] - backgrounds = test_data["backgrounds"] - - covars, _ = quat_scale_to_covar_preci(quats, scales, compute_preci=False, triu=True) + Ks, viewmats, height, width = ( + test_data["Ks"], + test_data["viewmats"], + test_data["height"], + test_data["width"], + ) + quats, scales, means, opacities = ( + test_data["quats"], + test_data["scales"] * 0.1, + test_data["means"], + test_data["opacities"], + ) + colors, backgrounds = test_data["colors"], test_data["backgrounds"] - # Project Gaussians to 2D - radii, means2d, depths, conics, compensations = fully_fused_projection( + covars, _ = gsplat.quat_scale_to_covar_preci( + quats, scales, compute_preci=False, triu=True + ) + radii, means2d, depths, conics, _ = gsplat.fully_fused_projection( means, covars, None, None, viewmats, Ks, width, height ) opacities = torch.broadcast_to(opacities[..., None, :], batch_dims + (C, N)) - # Identify intersecting tiles tile_size = 16 if channels <= 32 else 4 - tile_width = math.ceil(width / float(tile_size)) - tile_height = math.ceil(height / float(tile_size)) - tiles_per_gauss, isect_ids, flatten_ids = isect_tiles( + tile_width, tile_height = math.ceil(width / float(tile_size)), math.ceil( + height / float(tile_size) + ) + tiles_per_gauss, isect_ids, flatten_ids = gsplat.isect_tiles( means2d, radii, depths, tile_size, tile_width, tile_height ) - isect_offsets = isect_offset_encode(isect_ids, I, tile_width, tile_height) - isect_offsets = isect_offsets.reshape(batch_dims + (C, tile_height, tile_width)) + isect_offsets = gsplat.isect_offset_encode( + isect_ids, I, tile_width, tile_height + ).reshape(batch_dims + (C, tile_height, tile_width)) means2d.requires_grad = True conics.requires_grad = True @@ -549,8 +563,7 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. opacities.requires_grad = True backgrounds.requires_grad = True - # forward - render_colors, render_alphas = rasterize_to_pixels( + render_colors, render_alphas = gsplat.rasterize_to_pixels( means2d, conics, colors, @@ -562,80 +575,40 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. flatten_ids, backgrounds=backgrounds, ) - _render_colors, _render_alphas = _rasterize_to_pixels( - means2d, - conics, - colors, - opacities, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - backgrounds=backgrounds, - ) - torch.testing.assert_close(render_colors, _render_colors) - torch.testing.assert_close(render_alphas, _render_alphas) - # backward - v_render_colors = torch.randn_like(render_colors) - v_render_alphas = torch.randn_like(render_alphas) + if gsplat.BACKEND != "sycl": # nerfacc required for comparison + _render_colors, _render_alphas = _rasterize_to_pixels( + means2d, + conics, + colors, + opacities, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + backgrounds=backgrounds, + ) + torch.testing.assert_close(render_colors, _render_colors) + torch.testing.assert_close(render_alphas, _render_alphas) - v_means2d, v_conics, v_colors, v_opacities, v_backgrounds = torch.autograd.grad( + v_render_colors, v_render_alphas = torch.randn_like( + render_colors + ), torch.randn_like(render_alphas) + grads = torch.autograd.grad( (render_colors * v_render_colors).sum() + (render_alphas * v_render_alphas).sum(), (means2d, conics, colors, opacities, backgrounds), ) - ( - _v_means2d, - _v_conics, - _v_colors, - _v_opacities, - _v_backgrounds, - ) = torch.autograd.grad( - (_render_colors * v_render_colors).sum() - + (_render_alphas * v_render_alphas).sum(), - (means2d, conics, colors, opacities, backgrounds), - ) - torch.testing.assert_close(v_means2d, _v_means2d, rtol=5e-3, atol=5e-3) - torch.testing.assert_close(v_conics, _v_conics, rtol=1e-3, atol=1e-3) - torch.testing.assert_close(v_colors, _v_colors, rtol=1e-3, atol=1e-3) - torch.testing.assert_close(v_opacities, _v_opacities, rtol=8e-3, atol=6e-3) - torch.testing.assert_close(v_backgrounds, _v_backgrounds, rtol=1e-3, atol=1e-3) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") -@pytest.mark.parametrize("sh_degree", [0, 1, 2, 3, 4]) -@pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) -def test_sh(test_data, sh_degree: int, batch_dims: Tuple[int, ...]): - from gsplat.cuda._torch_impl import _spherical_harmonics - from gsplat.cuda._wrapper import spherical_harmonics - torch.manual_seed(42) - - N = 1000 - test_data = { - "coeffs": torch.randn(N, (4 + 1) ** 2, 3, device=device), - "dirs": torch.randn(N, 3, device=device), - } - test_data = expand(test_data, batch_dims) - coeffs = test_data["coeffs"] - dirs = test_data["dirs"] - coeffs.requires_grad = True - dirs.requires_grad = True - - colors = spherical_harmonics(sh_degree, dirs, coeffs) - _colors = _spherical_harmonics(sh_degree, dirs, coeffs) - torch.testing.assert_close(colors, _colors, rtol=1e-4, atol=1e-4) - - v_colors = torch.randn_like(colors) - - v_coeffs, v_dirs = torch.autograd.grad( - (colors * v_colors).sum(), (coeffs, dirs), retain_graph=True, allow_unused=True - ) - _v_coeffs, _v_dirs = torch.autograd.grad( - (_colors * v_colors).sum(), (coeffs, dirs), retain_graph=True, allow_unused=True - ) - torch.testing.assert_close(v_coeffs, _v_coeffs, rtol=1e-4, atol=1e-4) - if sh_degree > 0: - torch.testing.assert_close(v_dirs, _v_dirs, rtol=1e-4, atol=1e-4) + if gsplat.BACKEND != "sycl": # nerfacc required for comparison + _grads = torch.autograd.grad( + (_render_colors * v_render_colors).sum() + + (_render_alphas * v_render_alphas).sum(), + (means2d, conics, colors, opacities, backgrounds), + ) + torch.testing.assert_close(grads[0], _grads[0], rtol=5e-3, atol=5e-3) + torch.testing.assert_close(grads[1], _grads[1], rtol=1e-3, atol=1e-3) + torch.testing.assert_close(grads[2], _grads[2], rtol=1e-3, atol=1e-3) + torch.testing.assert_close(grads[3], _grads[3], rtol=8e-3, atol=6e-3) + torch.testing.assert_close(grads[4], _grads[4], rtol=1e-3, atol=1e-3) diff --git a/tests/test_rasterization.py b/tests/test_rasterization.py index 50247a0f..df53c5af 100644 --- a/tests/test_rasterization.py +++ b/tests/test_rasterization.py @@ -11,10 +11,23 @@ import pytest import torch -device = torch.device("cuda:0") +# device = torch.device("cuda:0") +import gsplat +if gsplat.BACKEND == "sycl": + device = torch.device("xpu:0") +elif gsplat.BACKEND == "cuda": + device = torch.device("cuda:0") +else: + device = None -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), + reason="No CUDA or SYCL XPU backend available", +) + + +@requires_backend @pytest.mark.parametrize("per_view_color", [True, False]) @pytest.mark.parametrize("sh_degree", [None, 3]) @pytest.mark.parametrize("render_mode", ["RGB", "RGB+D", "D"]) @@ -81,18 +94,19 @@ def test_rasterization( elif render_mode == "RGB+D": assert renders.shape == batch_dims + (C, height, width, 4) - _renders, _alphas, _meta = _rasterization( - means=means, - quats=quats, - scales=scales, - opacities=opacities, - colors=colors, - viewmats=viewmats, - Ks=Ks, - width=width, - height=height, - sh_degree=sh_degree, - render_mode=render_mode, - ) - torch.testing.assert_close(renders, _renders, rtol=1e-4, atol=1e-4) - torch.testing.assert_close(alphas, _alphas, rtol=1e-4, atol=1e-4) + if gsplat.BACKEND != "sycl": # nerfacc required for comparison + _renders, _alphas, _meta = _rasterization( + means=means, + quats=quats, + scales=scales, + opacities=opacities, + colors=colors, + viewmats=viewmats, + Ks=Ks, + width=width, + height=height, + sh_degree=sh_degree, + render_mode=render_mode, + ) + torch.testing.assert_close(renders, _renders, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(alphas, _alphas, rtol=1e-4, atol=1e-4) diff --git a/tests/test_strategy.py b/tests/test_strategy.py index 03115415..a5800d7d 100644 --- a/tests/test_strategy.py +++ b/tests/test_strategy.py @@ -8,11 +8,21 @@ import pytest import torch +import gsplat -device = torch.device("cuda:0") +if gsplat.BACKEND == "sycl": + device = torch.device("xpu:0") +elif gsplat.BACKEND == "cuda": + device = torch.device("cuda:0") +else: + device = torch.device("cpu") +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" +) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") + +@requires_backend def test_strategy(): from gsplat.rendering import rasterization from gsplat.strategy import DefaultStrategy, MCMCStrategy @@ -62,7 +72,7 @@ def test_strategy(): strategy.step_post_backward(params, optimizers, state, step=600, info=info, lr=1e-3) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend def test_strategy_requires_grad(): from gsplat.rendering import rasterization from gsplat.strategy import DefaultStrategy, MCMCStrategy