From 908956846156c162b8177013156ad0b15083e774 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Wed, 13 Aug 2025 15:09:17 +0000 Subject: [PATCH 01/56] added sycl code --- gsplat/__init__.py | 95 +- gsplat/{cuda => }/_torch_impl.py | 2 +- gsplat/{cuda => }/_torch_impl_2dgs.py | 4 +- gsplat/cuda/_wrapper.py | 2 +- gsplat/cuda/csrc/third_party/glm | 2 +- gsplat/strategy/ops.py | 13 +- gsplat/sycl/CMakeLists.txt | 119 + gsplat/sycl/__init__.py | 0 gsplat/sycl/_backend.py | 7 + gsplat/sycl/_wrapper.py | 2606 +++++++++++++++++ gsplat/sycl/ext.cpp | 104 + gsplat/sycl/include/Cameras.h | 58 + gsplat/sycl/include/Common.h | 32 + gsplat/sycl/include/Ops.h | 567 ++++ gsplat/sycl/include/gsplat_sycl_utils.hpp | 75 + gsplat/sycl/include/helpers.hpp | 14 + .../include/kernels/ComputeShBwdKernel.hpp | 69 + .../include/kernels/ComputeShFwdKernel.hpp | 55 + .../kernels/FullyFusedProjectionBwdKernel.hpp | 273 ++ .../kernels/FullyFusedProjectionFwdKernel.hpp | 220 ++ .../kernels/IsectOffsetEncodeKernel.hpp | 72 + .../sycl/include/kernels/IsectTilesKernel.hpp | 143 + gsplat/sycl/include/kernels/ProjBwdKernel.hpp | 140 + gsplat/sycl/include/kernels/ProjFwdKernel.hpp | 88 + .../QuatScaleToCovarPreciBwdKernel.hpp | 121 + .../QuatScaleToCovarPreciFwdKernel.hpp | 94 + .../kernels/RasterizeToPixelsBwdKernel.hpp | 370 +++ .../kernels/RasterizeToPixelsFwdKernel.hpp | 269 ++ .../include/kernels/WorldToCamBwdKernel.hpp | 124 + .../include/kernels/WorldToCamFwdKernel.hpp | 98 + gsplat/sycl/include/proj.hpp | 345 +++ gsplat/sycl/include/quat.hpp | 59 + .../include/quat_scale_to_covar_preci.hpp | 125 + gsplat/sycl/include/spherical_harmonics.hpp | 362 +++ gsplat/sycl/include/transform.hpp | 69 + gsplat/sycl/include/types.hpp | 22 + gsplat/sycl/include/utils.hpp | 84 + gsplat/sycl/src/adam.cpp | 23 + gsplat/sycl/src/intersect_offset.cpp | 58 + gsplat/sycl/src/intersect_tile.cpp | 121 + gsplat/sycl/src/null.cpp | 13 + gsplat/sycl/src/projection_2dgs_fused_bwd.cpp | 32 + gsplat/sycl/src/projection_2dgs_fused_fwd.cpp | 31 + .../sycl/src/projection_2dgs_packed_bwd.cpp | 35 + .../sycl/src/projection_2dgs_packed_fwd.cpp | 34 + .../src/projection_ewa_3dgs_fused_bwd.cpp | 36 + .../src/projection_ewa_3dgs_fused_fwd.cpp | 35 + .../src/projection_ewa_3dgs_packed_bwd.cpp | 39 + .../src/projection_ewa_3dgs_packed_fwd.cpp | 39 + gsplat/sycl/src/projection_ewa_simple_bwd.cpp | 70 + gsplat/sycl/src/projection_ewa_simple_fwd.cpp | 63 + gsplat/sycl/src/projection_ut_3dgs_fused.cpp | 43 + .../src/quat_scale_to_covar_preci_bwd.cpp | 60 + .../src/quat_scale_to_covar_preci_fwd.cpp | 65 + gsplat/sycl/src/rasterize_to_indices_2dgs.cpp | 28 + gsplat/sycl/src/rasterize_to_indices_3dgs.cpp | 28 + .../sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 51 + .../sycl/src/rasterize_to_pixels_2dgs_fwd.cpp | 37 + .../sycl/src/rasterize_to_pixels_3dgs_bwd.cpp | 195 ++ .../sycl/src/rasterize_to_pixels_3dgs_fwd.cpp | 154 + ...asterize_to_pixels_from_world_3dgs_bwd.cpp | 49 + ...asterize_to_pixels_from_world_3dgs_fwd.cpp | 43 + gsplat/sycl/src/relocation.cpp | 19 + gsplat/sycl/src/spherical_harmonics_bwd.cpp | 65 + gsplat/sycl/src/spherical_harmonics_fwd.cpp | 59 + setup.py | 80 +- tests/test_basic.py | 585 +--- 67 files changed, 8621 insertions(+), 471 deletions(-) rename gsplat/{cuda => }/_torch_impl.py (99%) rename gsplat/{cuda => }/_torch_impl_2dgs.py (99%) create mode 100644 gsplat/sycl/CMakeLists.txt create mode 100644 gsplat/sycl/__init__.py create mode 100644 gsplat/sycl/_backend.py create mode 100644 gsplat/sycl/_wrapper.py create mode 100644 gsplat/sycl/ext.cpp create mode 100644 gsplat/sycl/include/Cameras.h create mode 100644 gsplat/sycl/include/Common.h create mode 100644 gsplat/sycl/include/Ops.h create mode 100644 gsplat/sycl/include/gsplat_sycl_utils.hpp create mode 100644 gsplat/sycl/include/helpers.hpp create mode 100644 gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp create mode 100644 gsplat/sycl/include/kernels/IsectTilesKernel.hpp create mode 100644 gsplat/sycl/include/kernels/ProjBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/ProjFwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp create mode 100644 gsplat/sycl/include/proj.hpp create mode 100644 gsplat/sycl/include/quat.hpp create mode 100644 gsplat/sycl/include/quat_scale_to_covar_preci.hpp create mode 100644 gsplat/sycl/include/spherical_harmonics.hpp create mode 100644 gsplat/sycl/include/transform.hpp create mode 100644 gsplat/sycl/include/types.hpp create mode 100644 gsplat/sycl/include/utils.hpp create mode 100644 gsplat/sycl/src/adam.cpp create mode 100644 gsplat/sycl/src/intersect_offset.cpp create mode 100644 gsplat/sycl/src/intersect_tile.cpp create mode 100644 gsplat/sycl/src/null.cpp create mode 100644 gsplat/sycl/src/projection_2dgs_fused_bwd.cpp create mode 100644 gsplat/sycl/src/projection_2dgs_fused_fwd.cpp create mode 100644 gsplat/sycl/src/projection_2dgs_packed_bwd.cpp create mode 100644 gsplat/sycl/src/projection_2dgs_packed_fwd.cpp create mode 100644 gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp create mode 100644 gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp create mode 100644 gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp create mode 100644 gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp create mode 100644 gsplat/sycl/src/projection_ewa_simple_bwd.cpp create mode 100644 gsplat/sycl/src/projection_ewa_simple_fwd.cpp create mode 100644 gsplat/sycl/src/projection_ut_3dgs_fused.cpp create mode 100644 gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp create mode 100644 gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp create mode 100644 gsplat/sycl/src/rasterize_to_indices_2dgs.cpp create mode 100644 gsplat/sycl/src/rasterize_to_indices_3dgs.cpp create mode 100644 gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp create mode 100644 gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp create mode 100644 gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp create mode 100644 gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp create mode 100644 gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp create mode 100644 gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp create mode 100644 gsplat/sycl/src/relocation.cpp create mode 100644 gsplat/sycl/src/spherical_harmonics_bwd.cpp create mode 100644 gsplat/sycl/src/spherical_harmonics_fwd.cpp diff --git a/gsplat/__init__.py b/gsplat/__init__.py index 27c7b0a9..5cf505b0 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -1,25 +1,72 @@ +import os +import sys +import torch import warnings +BACKEND: str = "" + +FORCE_BACKEND = os.getenv("GSPLAT_BACKEND", "").lower() + +if FORCE_BACKEND == "cuda" or (FORCE_BACKEND == "" and torch.cuda.is_available()): + try: + BACKEND = "cuda" + from .cuda._wrapper import ( + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + proj, + quat_scale_to_covar_preci, + rasterize_to_indices_in_range, + rasterize_to_indices_in_range_2dgs, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + world_to_cam, + ) + print("gsplat: CUDA backend successfully loaded.", file=sys.stderr) + except ImportError: + if FORCE_BACKEND == "cuda": + print("gsplat: Error! GSPLAT_BACKEND=cuda was set but CUDA backend failed to load.", file=sys.stderr) + pass + +if not BACKEND and (FORCE_BACKEND == "sycl" or FORCE_BACKEND == ""): + try: + BACKEND = "sycl" + from .sycl._wrapper import ( + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + proj, + quat_scale_to_covar_preci, + rasterize_to_indices_in_range, + rasterize_to_indices_in_range_2dgs, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + world_to_cam, + ) + print("gsplat: SYCL backend successfully loaded.", file=sys.stderr) + except ImportError as e: + if FORCE_BACKEND == "sycl": + print(f"gsplat: Error! GSPLAT_BACKEND=sycl was set but SYCL backend failed to load: {e}", file=sys.stderr) + pass + +if not BACKEND: + print( + "gsplat: Warning! No high-performance backend (CUDA or SYCL) found.", + file=sys.stderr, + ) + + from .compression import PngCompression -from .cuda._torch_impl import accumulate -from .cuda._torch_impl_2dgs import accumulate_2dgs -from .cuda._wrapper import ( - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - proj, - quat_scale_to_covar_preci, - rasterize_to_indices_in_range, - rasterize_to_indices_in_range_2dgs, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - world_to_cam, -) from .exporter import export_splats from .optimizers import SelectiveAdam from .rendering import ( @@ -31,7 +78,9 @@ from .strategy import DefaultStrategy, MCMCStrategy, Strategy from .version import __version__ -all = [ + +__all__ = [ + "BACKEND", "PngCompression", "DefaultStrategy", "MCMCStrategy", @@ -47,16 +96,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 +] \ No newline at end of file diff --git a/gsplat/cuda/_torch_impl.py b/gsplat/_torch_impl.py similarity index 99% rename from gsplat/cuda/_torch_impl.py rename to gsplat/_torch_impl.py index 29888d3d..ab20ab5f 100644 --- a/gsplat/cuda/_torch_impl.py +++ b/gsplat/_torch_impl.py @@ -639,7 +639,7 @@ def _rasterize_to_pixels( This function requires the `nerfacc` package to be installed. Please install it using the following command `pip install nerfacc`. """ - from ._wrapper import rasterize_to_indices_in_range + from .cuda._wrapper import rasterize_to_indices_in_range image_dims = means2d.shape[:-2] channels = colors.shape[-1] diff --git a/gsplat/cuda/_torch_impl_2dgs.py b/gsplat/_torch_impl_2dgs.py similarity index 99% rename from gsplat/cuda/_torch_impl_2dgs.py rename to gsplat/_torch_impl_2dgs.py index 96ae8695..4f0fd4c3 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( @@ -231,7 +231,7 @@ def _rasterize_to_pixels_2dgs( This function requires the `nerfacc` package to be installed. Please install it using the following command `pip install nerfacc`. """ - from ._wrapper import rasterize_to_indices_in_range_2dgs + from .cuda._wrapper import rasterize_to_indices_in_range_2dgs image_dims = means2d.shape[:-2] channels = colors.shape[-1] diff --git a/gsplat/cuda/_wrapper.py b/gsplat/cuda/_wrapper.py index 286685ff..74b01ca5 100644 --- a/gsplat/cuda/_wrapper.py +++ b/gsplat/cuda/_wrapper.py @@ -112,7 +112,7 @@ def world_to_cam( - **Gaussian means in camera coordinate system**. [..., C, N, 3] - **Gaussian covariances in camera coordinate system**. [..., C, N, 3, 3] """ - from ._torch_impl import _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 " 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/strategy/ops.py b/gsplat/strategy/ops.py index 83c90a25..d14b37e2 100644 --- a/gsplat/strategy/ops.py +++ b/gsplat/strategy/ops.py @@ -5,7 +5,18 @@ import torch.nn.functional as F from torch import Tensor -from gsplat import quat_scale_to_covar_preci +from gsplat import BACKEND +if BACKEND == "cuda": + from gsplat.cuda._wrapper import ( + quat_scale_to_covar_preci, + ) +elif BACKEND == "sycl": + from gsplat.sycl._wrapper import ( + quat_scale_to_covar_preci, + ) +else: + raise ImportError("gsplat: No backend loaded, cannot import strategy ops.") + from gsplat.relocation import compute_relocation from gsplat.utils import normalized_quat_to_rotmat diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt new file mode 100644 index 00000000..0e2373e0 --- /dev/null +++ b/gsplat/sycl/CMakeLists.txt @@ -0,0 +1,119 @@ +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 torch.utils; print(torch.utils.cmake_prefix_path)" + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE Torch_DIR_From_Python + RESULT_VARIABLE _torch_path_result + ERROR_QUIET +) +if(NOT _torch_path_result EQUAL 0) + message(WARNING "Failed to get Torch CMake path from Python. " + "Make sure PyTorch is installed in the Python environment: ${Python_EXECUTABLE}") + set(Torch_DIR_From_Python "") +endif() + +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. " + "Please ensure PyTorch is installed or set CMAKE_PREFIX_PATH/Torch_DIR manually.") +endif() +execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], 'libtorch_python.so'))" + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE TORCH_PYTHON_LIB +) + +if (NOT EXISTS "${TORCH_PYTHON_LIB}") + message(FATAL_ERROR "Could not find libtorch_python.so at ${TORCH_PYTHON_LIB}. Please check your PyTorch installation.") +else() + message(STATUS "Found torch_python library at: ${TORCH_PYTHON_LIB}") +endif() + + +set(PYBIND11_FINDPYTHON ON) +find_package(pybind11 CONFIG REQUIRED) + +set( SYCL_SOURCES + ${CMAKE_CURRENT_SOURCE_DIR}/ext.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/adam.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/intersect_offset.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/intersect_tile.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/null.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_fused_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_fused_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_packed_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_packed_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_fused_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_fused_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_packed_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_packed_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_simple_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_simple_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ut_3dgs_fused.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/quat_scale_to_covar_preci_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/quat_scale_to_covar_preci_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_indices_2dgs.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_indices_3dgs.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_2dgs_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_2dgs_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_3dgs_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_3dgs_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/relocation.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/spherical_harmonics_bwd.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/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) + + +# 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..216999ca --- /dev/null +++ b/gsplat/sycl/_backend.py @@ -0,0 +1,7 @@ +_C = None + +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") \ No newline at end of file diff --git a/gsplat/sycl/_wrapper.py b/gsplat/sycl/_wrapper.py new file mode 100644 index 00000000..2172bf88 --- /dev/null +++ b/gsplat/sycl/_wrapper.py @@ -0,0 +1,2606 @@ +import math +import warnings +from dataclasses import dataclass +from enum import Enum +from typing import Any, Callable, Optional, Tuple + +import torch +from torch import Tensor +from typing_extensions import Literal + + +def _make_lazy_sycl_func(name: str) -> Callable: + """Creates a lazy-loading function for the SYCL backend.""" + def call_sycl(*args, **kwargs): + # pylint: disable=import-outside-toplevel + from ._backend import _C + return getattr(_C, name)(*args, **kwargs) + return call_sycl + +def _make_lazy_sycl_obj(name: str) -> Any: + """Creates a lazy-loading object accessor for the SYCL backend.""" + # pylint: disable=import-outside-toplevel + from ._backend import _C + + obj = _C + for name_split in name.split("."): + obj = getattr(obj, name_split) + return obj + + +class RollingShutterType(Enum): + ROLLING_TOP_TO_BOTTOM = 0 + ROLLING_LEFT_TO_RIGHT = 1 + ROLLING_BOTTOM_TO_TOP = 2 + ROLLING_RIGHT_TO_LEFT = 3 + GLOBAL = 4 + + def to_cpp(self) -> Any: + return _make_lazy_sycl_obj(f"ShutterType.{self.name}") + + +@dataclass +class UnscentedTransformParameters: + # Sigma point parameters (see Gustafsson and Hendeby 2012, Wan and van der Merwe 2000) + alpha: float = 0.1 + beta: float = 2.0 + kappa: float = 0.0 + # Parameters controlling validity of the unscented transform results. Default 0.1 + # is 10% margin. + in_image_margin_factor: float = 0.1 + # True: all sigma points must be valid + require_all_sigma_points_valid: bool = True + + def to_cpp(self) -> Any: + p = _make_lazy_sycl_obj("UnscentedTransformParameters")() + p.alpha = self.alpha + p.beta = self.beta + p.kappa = self.kappa + p.in_image_margin_factor = self.in_image_margin_factor + p.require_all_sigma_points_valid = self.require_all_sigma_points_valid + return p + + +@dataclass +class FThetaPolynomialType(Enum): + PIXELDIST_TO_ANGLE = 0 + ANGLE_TO_PIXELDIST = 1 + + def to_cpp(self) -> Any: + return _make_lazy_sycl_obj(f"FThetaPolynomialType.{self.name}") + + +@dataclass +class FThetaCameraDistortionParameters: + reference_poly: FThetaPolynomialType + pixeldist_to_angle_poly: Tuple[float, float, float, float, float, float] # [6] + angle_to_pixeldist_poly: Tuple[float, float, float, float, float, float] # [6] + max_angle: float + linear_cde: Tuple[float, float, float] # [3] + + def to_cpp(self) -> Any: + p = _make_lazy_sycl_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 + p.max_angle = self.max_angle + p.linear_cde = self.linear_cde + return p + + @classmethod + def to_cpp_default(cls) -> Any: + p = _make_lazy_sycl_obj("FThetaCameraDistortionParameters")() + return p + + +def world_to_cam( + means: Tensor, # [..., N, 3] + covars: Tensor, # [..., N, 3, 3] + viewmats: Tensor, # [..., C, 4, 4] +) -> Tuple[Tensor, Tensor]: + """Transforms Gaussians from world to camera coordinate system. + + Args: + means: Gaussian means. [..., N, 3] + covars: Gaussian covariances. [..., N, 3, 3] + viewmats: World-to-camera transformation matrices. [..., C, 4, 4] + + Returns: + A tuple: + + - **Gaussian means in camera coordinate system**. [..., C, N, 3] + - **Gaussian covariances in camera coordinate system**. [..., C, N, 3, 3] + """ + from .._torch_impl import _world_to_cam + + warnings.warn( + "world_to_cam() is removed from the sycl 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, + ) + batch_dims = means.shape[:-2] + N = means.shape[-2] + C = viewmats.shape[-3] + assert means.shape == batch_dims + (N, 3), means.shape + assert covars.shape == batch_dims + (N, 3, 3), covars.shape + assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape + means = means.contiguous() + covars = covars.contiguous() + viewmats = viewmats.contiguous() + return _world_to_cam(means, covars, viewmats) + + +def adam( + param: Tensor, + param_grad: Tensor, + exp_avg: Tensor, + exp_avg_sq: Tensor, + valid: Tensor, + lr: float, + b1: float, + b2: float, + eps: float, +) -> None: + _make_lazy_sycl_func("adam")( + param, param_grad, exp_avg, exp_avg_sq, valid, lr, b1, b2, eps + ) + + +def spherical_harmonics( + degrees_to_use: int, + dirs: Tensor, # [..., 3] + coeffs: Tensor, # [..., K, 3] + masks: Optional[Tensor] = None, # [...,] +) -> Tensor: + """Computes spherical harmonics. + + Args: + degrees_to_use: The degree to be used. + dirs: Directions. [..., 3] + coeffs: Coefficients. [..., K, 3] + masks: Optional boolen masks to skip some computation. [...,] Default: None. + + Returns: + Spherical harmonics. [..., 3] + """ + assert (degrees_to_use + 1) ** 2 <= coeffs.shape[-2], coeffs.shape + batch_dims = dirs.shape[:-1] + assert dirs.shape == batch_dims + (3,), dirs.shape + assert ( + (len(coeffs.shape) == len(batch_dims) + 2) + and coeffs.shape[:-2] == batch_dims + and coeffs.shape[-1] == 3 + ), coeffs.shape + if masks is not None: + assert masks.shape == batch_dims, masks.shape + masks = masks.contiguous() + return _SphericalHarmonics.apply( + degrees_to_use, dirs.contiguous(), coeffs.contiguous(), masks + ) + + +def quat_scale_to_covar_preci( + quats: Tensor, # [..., 4], + scales: Tensor, # [..., 3], + compute_covar: bool = True, + compute_preci: bool = True, + triu: bool = False, +) -> Tuple[Optional[Tensor], Optional[Tensor]]: + """Converts quaternions and scales to covariance and precision matrices. + + Args: + quats: Quaternions (No need to be normalized). [..., 4] + scales: Scales. [..., 3] + compute_covar: Whether to compute covariance matrices. Default: True. If False, + the returned covariance matrices will be None. + compute_preci: Whether to compute precision matrices. Default: True. If False, + the returned precision matrices will be None. + triu: If True, the return matrices will be upper triangular. Default: False. + + Returns: + A tuple: + + - **Covariance matrices**. If `triu` is True the returned shape is [..., 6], otherwise [..., 3, 3]. + - **Precision matrices**. If `triu` is True the returned shape is [..., 6], otherwise [..., 3, 3]. + """ + batch_dims = quats.shape[:-1] + assert quats.shape == batch_dims + (4,), quats.shape + assert scales.shape == batch_dims + (3,), scales.shape + quats = quats.contiguous() + scales = scales.contiguous() + covars, precis = _QuatScaleToCovarPreci.apply( + quats, scales, compute_covar, compute_preci, triu + ) + return covars if compute_covar else None, precis if compute_preci else None + + +def persp_proj( + means: Tensor, # [..., C, N, 3] + covars: Tensor, # [..., C, N, 3, 3] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, +) -> Tuple[Tensor, Tensor]: + """Perspective projection on Gaussians. + DEPRECATED: please use `proj` with `ortho=False` instead. + + Args: + means: Gaussian means. [..., C, N, 3] + covars: Gaussian covariances. [..., C, N, 3, 3] + Ks: Camera intrinsics. [..., C, 3, 3] + width: Image width. + height: Image height. + + Returns: + A tuple: + + - **Projected means**. [..., C, N, 2] + - **Projected covariances**. [..., C, N, 2, 2] + """ + warnings.warn( + "persp_proj is deprecated and will be removed in a future release. " + "Use proj with ortho=False instead.", + DeprecationWarning, + ) + return proj(means, covars, Ks, width, height, ortho=False) + + +def proj( + means: Tensor, # [..., C, N, 3] + covars: Tensor, # [..., C, N, 3, 3] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", +) -> Tuple[Tensor, Tensor]: + """Projection of Gaussians (perspective or orthographic). + + Args: + means: Gaussian means. [..., C, N, 3] + covars: Gaussian covariances. [..., C, N, 3, 3] + Ks: Camera intrinsics. [..., C, 3, 3] + width: Image width. + height: Image height. + + Returns: + A tuple: + + - **Projected means**. [..., C, N, 2] + - **Projected covariances**. [..., C, N, 2, 2] + """ + assert ( + camera_model != "ftheta" + ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" + + batch_dims = means.shape[:-3] + C, N = means.shape[-3:-1] + assert means.shape == batch_dims + (C, N, 3), means.shape + assert covars.shape == batch_dims + (C, N, 3, 3), covars.shape + assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape + means = means.contiguous() + covars = covars.contiguous() + Ks = Ks.contiguous() + return _Proj.apply(means, covars, Ks, width, height, camera_model) + + +def fully_fused_projection( + means: Tensor, # [..., N, 3] + covars: Optional[Tensor], # [..., N, 6] or None + quats: Optional[Tensor], # [..., N, 4] or None + scales: Optional[Tensor], # [..., N, 3] or None + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + eps2d: float = 0.3, + near_plane: float = 0.01, + far_plane: float = 1e10, + radius_clip: float = 0.0, + packed: bool = False, + sparse_grad: bool = False, + calc_compensations: bool = False, + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + opacities: Optional[Tensor] = None, # [..., N] or None +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Projects Gaussians to 2D. + + This function fuse the process of computing covariances + (:func:`quat_scale_to_covar_preci()`), transforming to camera space (:func:`world_to_cam()`), + and projection (:func:`proj()`). + + .. note:: + + During projection, we ignore the Gaussians that are outside of the camera frustum. + So not all the elements in the output tensors are valid. The output `radii` could serve as + 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 + 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 sycl kernel. Either `covars` or + {`quats`, `scales`} should be provided. + + Args: + means: Gaussian means. [..., N, 3] + covars: Gaussian covariances (flattened upper triangle). [..., N, 6] Optional. + quats: Quaternions (No need to be normalized). [..., N, 4] Optional. + scales: Scales. [..., N, 3] Optional. + viewmats: World-to-camera matrices. [..., C, 4, 4] + Ks: Camera intrinsics. [..., C, 3, 3] + width: Image width. + height: Image height. + eps2d: A epsilon added to the 2D covariance for numerical stability. Default: 0.3. + near_plane: Near plane distance. Default: 0.01. + far_plane: Far plane distance. Default: 1e10. + radius_clip: Gaussians with projected radii smaller than this value will be ignored. Default: 0.0. + packed: If True, the output tensors will be packed into a flattened tensor. Default: False. + sparse_grad: This is only effective when `packed` is True. If True, during backward the gradients + of {`means`, `covars`, `quats`, `scales`} will be a sparse Tensor in COO layout. Default: False. + calc_compensations: If True, a view-dependent opacity compensation factor will be computed, which + is useful for anti-aliasing. Default: False. + opacities: Gaussian opacities in range [0, 1]. If provided, will use it to compute a tighter bounds. + [..., N] or None. Default: None. + + Returns: + A tuple: + + If `packed` is True: + + - **batch_ids**. The batch indices of the projected Gaussians. Int32 tensor of shape [nnz]. + - **camera_ids**. The camera indices of the projected Gaussians. Int32 tensor of shape [nnz]. + - **gaussian_ids**. The column indices of the projected Gaussians. Int32 tensor of shape [nnz]. + - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [nnz, 2]. + - **means**. Projected Gaussian means in 2D. [nnz, 2] + - **depths**. The z-depth of the projected Gaussians. [nnz] + - **conics**. Inverse of the projected covariances. Return the flattend upper triangle with [nnz, 3] + - **compensations**. The view-dependent opacity compensation factor. [nnz] + + If `packed` is False: + + - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [..., C, N, 2]. + - **means**. Projected Gaussian means in 2D. [..., C, N, 2] + - **depths**. The z-depth of the projected Gaussians. [..., C, N] + - **conics**. Inverse of the projected covariances. Return the flattend upper triangle with [..., C, N, 3] + - **compensations**. The view-dependent opacity compensation factor. [..., C, N] + """ + batch_dims = means.shape[:-2] + N = means.shape[-2] + C = viewmats.shape[-3] + assert means.shape == batch_dims + (N, 3), means.shape + assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape + assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape + means = means.contiguous() + if covars is not None: + assert covars.shape == batch_dims + (N, 6), covars.shape + covars = covars.contiguous() + else: + assert quats is not None, "covars or quats is required" + assert scales is not None, "covars or scales is required" + assert quats.shape == batch_dims + (N, 4), quats.shape + assert scales.shape == batch_dims + (N, 3), scales.shape + quats = quats.contiguous() + scales = scales.contiguous() + if sparse_grad: + assert packed, "sparse_grad is only supported when packed is True" + assert batch_dims == (), "sparse_grad does not support batch dimensions" + if opacities is not None: + assert opacities.shape == batch_dims + (N,), opacities.shape + opacities = opacities.contiguous() + + assert ( + camera_model != "ftheta" + ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" + + viewmats = viewmats.contiguous() + Ks = Ks.contiguous() + if packed: + return _FullyFusedProjectionPacked.apply( + means, + covars, + quats, + scales, + viewmats, + Ks, + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + sparse_grad, + calc_compensations, + camera_model, + opacities, + ) + else: + return _FullyFusedProjection.apply( + means, + covars, + quats, + scales, + viewmats, + Ks, + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + calc_compensations, + camera_model, + opacities, + ) + + +@torch.no_grad() +def isect_tiles( + means2d: Tensor, # [..., N, 2] or [nnz, 2] + radii: Tensor, # [..., N, 2] or [nnz, 2] + depths: Tensor, # [..., N] or [nnz] + tile_size: int, + tile_width: int, + tile_height: int, + sort: bool = True, + segmented: bool = False, + packed: bool = False, + n_images: Optional[int] = None, + image_ids: Optional[Tensor] = None, + gaussian_ids: Optional[Tensor] = None, +) -> Tuple[Tensor, Tensor, Tensor]: + """Maps projected Gaussians to intersecting tiles. + + Args: + means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. + radii: Maximum radii of the projected Gaussians. [..., N, 2] if packed is False, [nnz, 2] if packed is True. + depths: Z-depth of the projected Gaussians. [..., N] if packed is False, [nnz] if packed is True. + tile_size: Tile size. + tile_width: Tile width. + tile_height: Tile height. + sort: If True, the returned intersections will be sorted by the intersection ids. Default: True. + segmented: If True, segmented radix sort will be used to sort the intersections. Default: False. + packed: If True, the input tensors are packed. Default: False. + n_images: Number of images. Required if packed is True. + image_ids: The image indices of the projected Gaussians. Required if packed is True. + gaussian_ids: The column indices of the projected Gaussians. Required if packed is True. + + Returns: + A tuple: + + - **Tiles per Gaussian**. The number of tiles intersected by each Gaussian. + Int32 [..., N] if packed is False, Int32 [nnz] if packed is True. + - **Intersection ids**. Each id is an 64-bit integer with the following + information: image_id (Xc bits) | tile_id (Xt bits) | depth (32 bits). + Xc and Xt are the maximum number of bits required to represent the image and + tile ids, respectively. Int64 [n_isects] + - **Flatten ids**. The global flatten indices in [I * N] or [nnz] (packed). [n_isects] + """ + if packed: + nnz = means2d.size(0) + assert means2d.shape == (nnz, 2), means2d.shape + assert radii.shape == (nnz, 2), radii.shape + assert depths.shape == (nnz,), depths.shape + assert image_ids is not None, "image_ids is required if packed is True" + assert gaussian_ids is not None, "gaussian_ids is required if packed is True" + assert n_images is not None, "n_images is required if packed is True" + image_ids = image_ids.contiguous() + gaussian_ids = gaussian_ids.contiguous() + I = n_images + + else: + image_dims = means2d.shape[:-2] + I = math.prod(image_dims) + N = means2d.shape[-2] + assert means2d.shape == image_dims + (N, 2), means2d.shape + 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_sycl_func("intersect_tile")( + means2d.contiguous(), + radii.contiguous(), + depths.contiguous(), + image_ids, + gaussian_ids, + I, + tile_size, + tile_width, + tile_height, + sort, + segmented, + ) + return tiles_per_gauss, isect_ids, flatten_ids + + +@torch.no_grad() +def isect_offset_encode( + isect_ids: Tensor, + n_images: int, + tile_width: int, + tile_height: int, +) -> Tensor: + """Encodes intersection ids to offsets. + + Args: + isect_ids: Intersection ids. [n_isects] + n_images: Number of images. + tile_width: Tile width. + tile_height: Tile height. + + Returns: + Offsets. [I, tile_height, tile_width] + """ + return _make_lazy_sycl_func("intersect_offset")( + isect_ids.contiguous(), n_images, tile_width, tile_height + ) + + +def rasterize_to_pixels( + means2d: Tensor, # [..., N, 2] or [nnz, 2] + conics: Tensor, # [..., N, 3] or [nnz, 3] + colors: Tensor, # [..., N, channels] or [nnz, channels] + opacities: Tensor, # [..., N] or [nnz] + image_width: int, + image_height: int, + tile_size: int, + isect_offsets: Tensor, # [..., tile_height, tile_width] + flatten_ids: Tensor, # [n_isects] + backgrounds: Optional[Tensor] = None, # [..., channels] + masks: Optional[Tensor] = None, # [..., tile_height, tile_width] + packed: bool = False, + absgrad: bool = False, +) -> Tuple[Tensor, Tensor]: + """Rasterizes Gaussians to pixels. + + Args: + means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. + conics: Inverse of the projected covariances with only upper triangle values. [..., N, 3] if packed is False, [nnz, 3] if packed is True. + colors: Gaussian colors or ND features. [..., N, channels] if packed is False, [nnz, channels] if packed is True. + opacities: Gaussian opacities that support per-view values. [..., N] if packed is False, [nnz] if packed is True. + image_width: Image width. + image_height: Image height. + tile_size: Tile size. + isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] + flatten_ids: The global flatten indices in [I * N] or [nnz] from `isect_tiles()`. [n_isects] + backgrounds: Background colors. [..., channels]. Default: None. + masks: Optional tile mask to skip rendering GS to masked tiles. [..., tile_height, tile_width]. Default: None. + packed: If True, the input tensors are expected to be packed with shape [nnz, ...]. Default: False. + absgrad: If True, the backward pass will compute a `.absgrad` attribute for `means2d`. Default: False. + + Returns: + A tuple: + + - **Rendered colors**. [..., image_height, image_width, channels] + - **Rendered alphas**. [..., image_height, image_width, 1] + """ + + image_dims = means2d.shape[:-2] + channels = colors.shape[-1] + device = means2d.device + if packed: + nnz = means2d.size(0) + assert means2d.shape == (nnz, 2), means2d.shape + assert conics.shape == (nnz, 3), conics.shape + assert colors.shape[0] == nnz, colors.shape + assert opacities.shape == (nnz,), opacities.shape + else: + N = means2d.size(-2) + assert means2d.shape == image_dims + (N, 2), means2d.shape + assert conics.shape == image_dims + (N, 3), conics.shape + assert colors.shape == image_dims + (N, channels), colors.shape + assert opacities.shape == image_dims + (N,), opacities.shape + if backgrounds is not None: + assert backgrounds.shape == image_dims + (channels,), backgrounds.shape + backgrounds = backgrounds.contiguous() + if masks is not None: + assert masks.shape == isect_offsets.shape, masks.shape + masks = masks.contiguous() + + # Pad the channels to the nearest supported number if necessary + if channels > 513 or channels == 0: + # TODO: maybe worth to support zero channels? + raise ValueError(f"Unsupported number of color channels: {channels}") + if channels not in ( + 1, + 2, + 3, + 4, + 5, + 8, + 9, + 16, + 17, + 32, + 33, + 64, + 65, + 128, + 129, + 256, + 257, + 512, + 513, + ): + padded_channels = (1 << (channels - 1).bit_length()) - channels + colors = torch.cat( + [ + colors, + torch.zeros(*colors.shape[:-1], padded_channels, device=device), + ], + dim=-1, + ) + if backgrounds is not None: + backgrounds = torch.cat( + [ + backgrounds, + torch.zeros( + *backgrounds.shape[:-1], padded_channels, device=device + ), + ], + dim=-1, + ) + else: + padded_channels = 0 + + tile_height, tile_width = isect_offsets.shape[-2:] + assert ( + tile_height * tile_size >= image_height + ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" + assert ( + tile_width * tile_size >= image_width + ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" + + render_colors, render_alphas = _RasterizeToPixels.apply( + means2d.contiguous(), + conics.contiguous(), + colors.contiguous(), + opacities.contiguous(), + backgrounds, + masks, + image_width, + image_height, + tile_size, + isect_offsets.contiguous(), + flatten_ids.contiguous(), + absgrad, + ) + + if padded_channels > 0: + render_colors = render_colors[..., :-padded_channels] + return render_colors, render_alphas + + +def rasterize_to_pixels_eval3d( + means: Tensor, # [..., N, 3] + quats: Tensor, # [..., N, 4] + scales: Tensor, # [..., N, 3] + colors: Tensor, # [..., C, N, channels] or [nnz, channels] + opacities: Tensor, # [..., C, N] or [nnz] + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + image_width: int, + image_height: int, + tile_size: int, + isect_offsets: Tensor, # [..., C, tile_height, tile_width] + flatten_ids: Tensor, # [n_isects] + backgrounds: Optional[Tensor] = None, # [..., C, channels] + masks: Optional[Tensor] = None, # [..., C, tile_height, tile_width] + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + ut_params: UnscentedTransformParameters = UnscentedTransformParameters(), + # distortion + 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, + # rolling shutter + rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, + viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] +) -> Tuple[Tensor, Tensor]: + """Rasterizes Gaussians to pixels. + + Similar to `rasterize_to_pixels()`, but compute the Gaussian responses in the + 3D world space instead of the 2D image space. Supports rolling shutter and + camera distortion. + + Returns: + A tuple: + + - **Rendered colors**. [..., C, image_height, image_width, channels] + - **Rendered alphas**. [..., C, image_height, image_width, 1] + """ + batch_dims = means.shape[:-2] + num_batch_dims = len(batch_dims) + N = means.size(-2) + C = viewmats.size(-3) + channels = colors.shape[-1] + device = means.device + + assert means.shape == batch_dims + (N, 3), means.shape + assert quats.shape == batch_dims + (N, 4), quats.shape + assert scales.shape == batch_dims + (N, 3), scales.shape + assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape + assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape + + assert colors.ndim in (num_batch_dims + 2, num_batch_dims + 3), colors.shape + if colors.ndim == num_batch_dims + 2: + raise NotImplementedError("packed mode is not supported yet") + assert ( + colors.shape[:-2] == batch_dims and colors.shape[-1] == channels + ), colors.shape + else: + assert colors.shape == batch_dims + (C, N, channels), colors.shape + assert opacities.shape == colors.shape[:-1], opacities.shape + + if backgrounds is not None: + assert backgrounds.shape == batch_dims + (C, channels), backgrounds.shape + backgrounds = backgrounds.contiguous() + + if masks is not None: + assert masks.shape == isect_offsets.shape, masks.shape + masks = masks.contiguous() + + if radial_coeffs is not None: + assert radial_coeffs.shape[:-1] == batch_dims + (C,) and radial_coeffs.shape[ + -1 + ] in (6, 4), radial_coeffs.shape + radial_coeffs = radial_coeffs.contiguous() + + if tangential_coeffs is not None: + assert tangential_coeffs.shape == batch_dims + (C, 2), tangential_coeffs.shape + tangential_coeffs = tangential_coeffs.contiguous() + + if thin_prism_coeffs is not None: + assert thin_prism_coeffs.shape == batch_dims + (C, 4), thin_prism_coeffs.shape + thin_prism_coeffs = thin_prism_coeffs.contiguous() + + if viewmats_rs is not None: + assert viewmats_rs.shape == batch_dims + (C, 4, 4), viewmats_rs.shape + viewmats_rs = viewmats_rs.contiguous() + + # Pad the channels to the nearest supported number if necessary + channels = colors.shape[-1] + if channels > 513 or channels == 0: + # TODO: maybe worth to support zero channels? + raise ValueError(f"Unsupported number of color channels: {channels}") + if channels not in ( + 1, + 2, + 3, + 4, + 5, + 8, + 9, + 16, + 17, + 32, + 33, + 64, + 65, + 128, + 129, + 256, + 257, + 512, + 513, + ): + padded_channels = (1 << (channels - 1).bit_length()) - channels + colors = torch.cat( + [ + colors, + torch.zeros(*colors.shape[:-1], padded_channels, device=device), + ], + dim=-1, + ) + if backgrounds is not None: + backgrounds = torch.cat( + [ + backgrounds, + torch.zeros( + *backgrounds.shape[:-1], padded_channels, device=device + ), + ], + dim=-1, + ) + else: + padded_channels = 0 + + tile_height, tile_width = isect_offsets.shape[-2:] + assert ( + tile_height * tile_size >= image_height + ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" + assert ( + tile_width * tile_size >= image_width + ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" + + render_colors, render_alphas = _RasterizeToPixelsEval3D.apply( + means.contiguous(), + quats.contiguous(), + scales.contiguous(), + colors.contiguous(), + opacities.contiguous(), + backgrounds.contiguous() if backgrounds is not None else None, + masks.contiguous() if masks is not None else None, + viewmats.contiguous(), + Ks.contiguous(), + image_width, + image_height, + tile_size, + isect_offsets.contiguous(), + flatten_ids.contiguous(), + camera_model, + ut_params, + # distortion + 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, + # rolling shutter + rolling_shutter, + viewmats_rs.contiguous() if viewmats_rs is not None else None, + ) + + if padded_channels > 0: + render_colors = render_colors[..., :-padded_channels] + return render_colors, render_alphas + + +@torch.no_grad() +def rasterize_to_indices_in_range( + range_start: int, + range_end: int, + transmittances: Tensor, # [..., image_height, image_width] + means2d: Tensor, # [..., N, 2] + conics: Tensor, # [..., N, 3] + opacities: Tensor, # [..., N] + image_width: int, + image_height: int, + tile_size: int, + isect_offsets: Tensor, # [..., tile_height, tile_width] + flatten_ids: Tensor, # [n_isects] +) -> Tuple[Tensor, Tensor, Tensor]: + """Rasterizes a batch of Gaussians to images but only returns the indices. + + .. note:: + + This function supports iterative rasterization, in which each call of this function + will rasterize a batch of Gaussians from near to far, defined by `[range_start, range_end)`. + If a one-step full rasterization is desired, set `range_start` to 0 and `range_end` to a really + large number, e.g, 1e10. + + Args: + range_start: The start batch of Gaussians to be rasterized (inclusive). + range_end: The end batch of Gaussians to be rasterized (exclusive). + transmittances: Currently transmittances. [..., image_height, image_width] + means2d: Projected Gaussian means. [..., N, 2] + conics: Inverse of the projected covariances with only upper triangle values. [..., N, 3] + opacities: Gaussian opacities that support per-view values. [..., N] + image_width: Image width. + image_height: Image height. + tile_size: Tile size. + isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] + flatten_ids: The global flatten indices in [I * N] from `isect_tiles()`. [n_isects] + + Returns: + A tuple: + + - **Gaussian ids**. Gaussian ids for the pixel intersection. A flattened list of shape [M]. + - **Pixel ids**. pixel indices (row-major). A flattened list of shape [M]. + - **Image ids**. image indices. A flattened list of shape [M]. + """ + + image_dims = means2d.shape[:-2] + tile_height, tile_width = isect_offsets.shape[-2:] + N = means2d.shape[-2] + assert transmittances.shape == image_dims + ( + image_height, + image_width, + ), transmittances.shape + assert means2d.shape == image_dims + (N, 2), means2d.shape + assert conics.shape == image_dims + (N, 3), conics.shape + assert opacities.shape == image_dims + (N,), opacities.shape + assert isect_offsets.shape == image_dims + ( + tile_height, + tile_width, + ), isect_offsets.shape + assert ( + tile_height * tile_size >= image_height + ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" + assert ( + tile_width * tile_size >= image_width + ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" + + out_gauss_ids, out_indices = _make_lazy_sycl_func("rasterize_to_indices_3dgs")( + range_start, + range_end, + transmittances.contiguous(), + means2d.contiguous(), + conics.contiguous(), + opacities.contiguous(), + image_width, + image_height, + tile_size, + isect_offsets.contiguous(), + flatten_ids.contiguous(), + ) + out_pixel_ids = out_indices % (image_width * image_height) + out_image_ids = out_indices // (image_width * image_height) + return out_gauss_ids, out_pixel_ids, out_image_ids + + +class _QuatScaleToCovarPreci(torch.autograd.Function): + """Converts quaternions and scales to covariance and precision matrices.""" + + @staticmethod + def forward( + ctx, + quats: Tensor, # [..., 4], + scales: Tensor, # [..., 3], + compute_covar: bool = True, + compute_preci: bool = True, + triu: bool = False, + ) -> Tuple[Tensor, Tensor]: + covars, precis = _make_lazy_sycl_func("quat_scale_to_covar_preci_fwd")( + quats, scales, compute_covar, compute_preci, triu + ) + ctx.save_for_backward(quats, scales) + ctx.compute_covar = compute_covar + ctx.compute_preci = compute_preci + ctx.triu = triu + return covars, precis + + @staticmethod + def backward(ctx, v_covars: Tensor, v_precis: Tensor): + quats, scales = ctx.saved_tensors + compute_covar = ctx.compute_covar + compute_preci = ctx.compute_preci + triu = ctx.triu + if compute_covar and v_covars.is_sparse: + 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_sycl_func("quat_scale_to_covar_preci_bwd")( + quats, + scales, + triu, + v_covars.contiguous() if compute_covar else None, + v_precis.contiguous() if compute_preci else None, + ) + return v_quats, v_scales, None, None, None + + +class _Proj(torch.autograd.Function): + """Perspective fully_fused_projection on Gaussians.""" + + @staticmethod + def forward( + ctx, + means: Tensor, # [..., C, N, 3] + covars: Tensor, # [..., C, N, 3, 3] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + ) -> Tuple[Tensor, Tensor]: + assert ( + camera_model != "ftheta" + ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" + + camera_model_type = _make_lazy_sycl_obj( + f"CameraModelType.{camera_model.upper()}" + ) + + means2d, covars2d = _make_lazy_sycl_func("projection_ewa_simple_fwd")( + means, + covars, + Ks, + width, + height, + camera_model_type, + ) + ctx.save_for_backward(means, covars, Ks) + ctx.width = width + ctx.height = height + ctx.camera_model_type = camera_model_type + return means2d, covars2d + + @staticmethod + def backward(ctx, v_means2d: Tensor, v_covars2d: Tensor): + means, covars, Ks = ctx.saved_tensors + width = ctx.width + height = ctx.height + camera_model_type = ctx.camera_model_type + v_means, v_covars = _make_lazy_sycl_func("projection_ewa_simple_bwd")( + means, + covars, + Ks, + width, + height, + camera_model_type, + v_means2d.contiguous(), + v_covars2d.contiguous(), + ) + return v_means, v_covars, None, None, None, None + + +class _FullyFusedProjection(torch.autograd.Function): + """Projects Gaussians to 2D.""" + + @staticmethod + def forward( + ctx, + means: Tensor, # [..., N, 3] + covars: Tensor, # [..., N, 6] or None + quats: Tensor, # [..., N, 4] or None + scales: Tensor, # [..., N, 3] or None + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + eps2d: float, + near_plane: float, + far_plane: float, + radius_clip: float, + calc_compensations: bool, + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + opacities: Optional[Tensor] = None, # [..., N] or None + ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + assert ( + camera_model != "ftheta" + ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" + + camera_model_type = _make_lazy_sycl_obj( + f"CameraModelType.{camera_model.upper()}" + ) + + # "covars" and {"quats", "scales"} are mutually exclusive + radii, means2d, depths, conics, compensations = _make_lazy_sycl_func( + "projection_ewa_3dgs_fused_fwd" + )( + means, + covars, + quats, + scales, + opacities, + viewmats, + Ks, + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + calc_compensations, + camera_model_type, + ) + if not calc_compensations: + compensations = None + ctx.save_for_backward( + means, covars, quats, scales, viewmats, Ks, radii, conics, compensations + ) + ctx.width = width + ctx.height = height + ctx.eps2d = eps2d + ctx.camera_model_type = camera_model_type + + return radii, means2d, depths, conics, compensations + + @staticmethod + def backward(ctx, v_radii, v_means2d, v_depths, v_conics, v_compensations): + ( + means, + covars, + quats, + scales, + viewmats, + Ks, + radii, + conics, + compensations, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + eps2d = ctx.eps2d + 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_sycl_func( + "projection_ewa_3dgs_fused_bwd" + )( + means, + covars, + quats, + scales, + viewmats, + Ks, + width, + height, + eps2d, + camera_model_type, + radii, + conics, + compensations, + v_means2d.contiguous(), + v_depths.contiguous(), + v_conics.contiguous(), + v_compensations, + ctx.needs_input_grad[4], # viewmats_requires_grad + ) + if not ctx.needs_input_grad[0]: + v_means = None + if not ctx.needs_input_grad[1]: + v_covars = None + if not ctx.needs_input_grad[2]: + v_quats = None + if not ctx.needs_input_grad[3]: + v_scales = None + if not ctx.needs_input_grad[4]: + v_viewmats = None + return ( + v_means, + v_covars, + v_quats, + v_scales, + v_viewmats, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def fully_fused_projection_with_ut( + means: Tensor, # [..., N, 3] + quats: Tensor, # [..., N, 4] + scales: Tensor, # [..., N, 3] + opacities: Optional[Tensor], # [..., N] + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + eps2d: float = 0.3, + near_plane: float = 0.01, + far_plane: float = 1e10, + radius_clip: float = 0.0, + calc_compensations: bool = False, + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + ut_params: UnscentedTransformParameters = UnscentedTransformParameters(), + # distortion + 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, + # rolling shutter + rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, + viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] +) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + """Projects Gaussians to 2D using Unscented Transform (UT). + + similar to `fully_fused_projection()`, but supports camera distortion and + rolling shutter. + + .. warning:: + This function is not differentiable to any input. + """ + batch_dims = means.shape[:-2] + N = means.shape[-2] + C = viewmats.shape[-3] + assert means.shape == batch_dims + (N, 3), means.shape + assert quats.shape == batch_dims + (N, 4), quats.shape + assert scales.shape == batch_dims + (N, 3), scales.shape + if opacities is not None: + assert opacities.shape == batch_dims + (N,), opacities.shape + assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape + assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape + if radial_coeffs is not None: + assert radial_coeffs.shape[:-1] == batch_dims + (C,) and radial_coeffs.shape[ + -1 + ] in [6, 4], radial_coeffs.shape + if tangential_coeffs is not None: + assert tangential_coeffs.shape == batch_dims + (C, 2), tangential_coeffs.shape + if thin_prism_coeffs is not None: + assert thin_prism_coeffs.shape == batch_dims + (C, 4), thin_prism_coeffs.shape + if viewmats_rs is not None: + assert viewmats_rs.shape == batch_dims + (C, 4, 4), viewmats_rs.shape + + camera_model_type = _make_lazy_sycl_obj(f"CameraModelType.{camera_model.upper()}") + + radii, means2d, depths, conics, compensations = _make_lazy_sycl_func( + "projection_ut_3dgs_fused" + )( + means.contiguous(), + quats.contiguous(), + scales.contiguous(), + opacities.contiguous() if opacities is not None else None, + viewmats.contiguous(), + viewmats_rs.contiguous() if viewmats_rs is not None else None, + Ks.contiguous(), + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + calc_compensations, + camera_model_type, + ut_params.to_cpp(), + rolling_shutter.to_cpp(), + 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(), + ) + if not calc_compensations: + compensations = None + return radii, means2d, depths, conics, compensations + + +class _RasterizeToPixels(torch.autograd.Function): + """Rasterize gaussians""" + + @staticmethod + def forward( + ctx, + means2d: Tensor, # [..., N, 2] or [nnz, 2] + conics: Tensor, # [..., N, 3] or [nnz, 3] + colors: Tensor, # [..., N, channels] or [nnz, channels] + opacities: Tensor, # [..., N] or [nnz] + backgrounds: Tensor, # [..., channels], Optional + masks: Tensor, # [..., tile_height, tile_width], Optional + width: int, + height: int, + tile_size: int, + isect_offsets: Tensor, # [..., tile_height, tile_width] + flatten_ids: Tensor, # [n_isects] + absgrad: bool, + ) -> Tuple[Tensor, Tensor]: + render_colors, render_alphas, last_ids = _make_lazy_sycl_func( + "rasterize_to_pixels_3dgs_fwd" + )( + means2d, + conics, + colors, + opacities, + backgrounds, + masks, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + ) + + ctx.save_for_backward( + means2d, + conics, + colors, + opacities, + backgrounds, + masks, + isect_offsets, + flatten_ids, + render_alphas, + last_ids, + ) + ctx.width = width + ctx.height = height + ctx.tile_size = tile_size + ctx.absgrad = absgrad + + # double to float + render_alphas = render_alphas.float() + return render_colors, render_alphas + + @staticmethod + def backward( + ctx, + v_render_colors: Tensor, # [..., H, W, 3] + v_render_alphas: Tensor, # [..., H, W, 1] + ): + ( + means2d, + conics, + colors, + opacities, + backgrounds, + masks, + isect_offsets, + flatten_ids, + render_alphas, + last_ids, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + tile_size = ctx.tile_size + absgrad = ctx.absgrad + + ( + v_means2d_abs, + v_means2d, + v_conics, + v_colors, + v_opacities, + ) = _make_lazy_sycl_func("rasterize_to_pixels_3dgs_bwd")( + means2d, + conics, + colors, + opacities, + backgrounds, + masks, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + render_alphas, + last_ids, + v_render_colors.contiguous(), + v_render_alphas.contiguous(), + absgrad, + ) + + if absgrad: + means2d.absgrad = v_means2d_abs + + if ctx.needs_input_grad[4]: + v_backgrounds = (v_render_colors * (1.0 - render_alphas).float()).sum( + dim=(-3, -2) + ) + else: + v_backgrounds = None + + return ( + v_means2d, + v_conics, + v_colors, + v_opacities, + v_backgrounds, + None, + None, + None, + None, + None, + None, + None, + ) + + +class _RasterizeToPixelsEval3D(torch.autograd.Function): + """Rasterize gaussians""" + + @staticmethod + def forward( + ctx, + means: Tensor, # [..., N, 3] + quats: Tensor, # [..., N, 4] + scales: Tensor, # [..., N, 3] + colors: Tensor, # [..., C, N, D] or [nnz, D] + opacities: Tensor, # [..., C, N] or [nnz] + backgrounds: Tensor, # [..., C, D], Optional + masks: Tensor, # [..., C, tile_height, tile_width], Optional + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + tile_size: int, + isect_offsets: Tensor, # [..., C, tile_height, tile_width] + flatten_ids: Tensor, # [..., n_isects] + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + ut_params: UnscentedTransformParameters = UnscentedTransformParameters(), + # distortion + 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, + # rolling shutter + rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, + viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] + ) -> Tuple[Tensor, Tensor]: + ut_params = ut_params.to_cpp() + rs_type = rolling_shutter.to_cpp() + camera_model_type = _make_lazy_sycl_obj( + f"CameraModelType.{camera_model.upper()}" + ) + ftheta_coeffs = ( + ftheta_coeffs.to_cpp() + if ftheta_coeffs is not None + else FThetaCameraDistortionParameters.to_cpp_default() + ) + + render_colors, render_alphas, last_ids = _make_lazy_sycl_func( + "rasterize_to_pixels_from_world_3dgs_fwd" + )( + means, + quats, + scales, + colors, + opacities, + backgrounds, + masks, + width, + height, + tile_size, + viewmats, + viewmats_rs, + Ks, + camera_model_type, + ut_params, + rs_type, + radial_coeffs, + tangential_coeffs, + thin_prism_coeffs, + ftheta_coeffs, + isect_offsets, + flatten_ids, + ) + + ctx.save_for_backward( + means, + quats, + scales, + colors, + opacities, + backgrounds, + masks, + viewmats, + viewmats_rs, + Ks, + radial_coeffs, + tangential_coeffs, + thin_prism_coeffs, + isect_offsets, + flatten_ids, + render_alphas, + last_ids, + ) + ctx.width = width + ctx.height = height + ctx.ut_params = ut_params + ctx.rs_type = rs_type + ctx.camera_model_type = camera_model_type + ctx.tile_size = tile_size + ctx.ftheta_coeffs = ftheta_coeffs + + return render_colors, render_alphas + + @staticmethod + def backward( + ctx, + v_render_colors: Tensor, # [..., C, H, W, 3] + v_render_alphas: Tensor, # [..., C, H, W, 1] + ): + ( + means, + quats, + scales, + colors, + opacities, + backgrounds, + masks, + viewmats, + viewmats_rs, + Ks, + radial_coeffs, + tangential_coeffs, + thin_prism_coeffs, + isect_offsets, + flatten_ids, + render_alphas, + last_ids, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + ut_params = ctx.ut_params + rs_type = ctx.rs_type + camera_model_type = ctx.camera_model_type + tile_size = ctx.tile_size + ftheta_coeffs = ctx.ftheta_coeffs + + (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_sycl_func( + "rasterize_to_pixels_from_world_3dgs_bwd" + )( + means, + quats, + scales, + colors, + opacities, + backgrounds, + masks, + width, + height, + tile_size, + viewmats, + viewmats_rs, + Ks, + camera_model_type, + ut_params, + rs_type, + radial_coeffs, + tangential_coeffs, + thin_prism_coeffs, + ftheta_coeffs, + isect_offsets, + flatten_ids, + render_alphas, + last_ids, + v_render_colors.contiguous(), + v_render_alphas.contiguous(), + ) + + if ctx.needs_input_grad[5]: # backgrounds + v_backgrounds = (v_render_colors * (1.0 - render_alphas).float()).sum( + dim=(-3, -2) + ) + else: + v_backgrounds = None + + if ctx.needs_input_grad[7]: # viewmats + raise NotImplementedError + + return ( + v_means, + v_quats, + v_scales, + v_colors, + v_opacities, + v_backgrounds, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +class _FullyFusedProjectionPacked(torch.autograd.Function): + """Projects Gaussians to 2D. Return packed tensors.""" + + @staticmethod + def forward( + ctx, + means: Tensor, # [..., N, 3] + covars: Tensor, # [..., N, 6] or None + quats: Tensor, # [..., N, 4] or None + scales: Tensor, # [..., N, 3] or None + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + eps2d: float, + near_plane: float, + far_plane: float, + radius_clip: float, + sparse_grad: bool, + calc_compensations: bool, + camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", + opacities: Optional[Tensor] = None, # [..., N] or None + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + assert ( + camera_model != "ftheta" + ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" + + camera_model_type = _make_lazy_sycl_obj( + f"CameraModelType.{camera_model.upper()}" + ) + + ( + indptr, + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + compensations, + ) = _make_lazy_sycl_func("projection_ewa_3dgs_packed_fwd")( + means, + covars, # optional + quats, # optional + scales, # optional + opacities, # optional + viewmats, + Ks, + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + calc_compensations, + camera_model_type, + ) + if not calc_compensations: + compensations = None + ctx.save_for_backward( + batch_ids, + camera_ids, + gaussian_ids, + means, + covars, + quats, + scales, + viewmats, + Ks, + conics, + compensations, + ) + ctx.width = width + ctx.height = height + ctx.eps2d = eps2d + ctx.sparse_grad = sparse_grad + ctx.camera_model_type = camera_model_type + + return ( + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + compensations, + ) + + @staticmethod + def backward( + ctx, + v_batch_ids, + v_camera_ids, + v_gaussian_ids, + v_radii, + v_means2d, + v_depths, + v_conics, + v_compensations, + ): + ( + batch_ids, + camera_ids, + gaussian_ids, + means, + covars, + quats, + scales, + viewmats, + Ks, + conics, + compensations, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + eps2d = ctx.eps2d + sparse_grad = ctx.sparse_grad + 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_sycl_func( + "projection_ewa_3dgs_packed_bwd" + )( + means, + covars, + quats, + scales, + viewmats, + Ks, + width, + height, + eps2d, + camera_model_type, + batch_ids, + camera_ids, + gaussian_ids, + conics, + compensations, + v_means2d.contiguous(), + v_depths.contiguous(), + v_conics.contiguous(), + v_compensations, + ctx.needs_input_grad[4], # viewmats_requires_grad + sparse_grad, + ) + + if sparse_grad: + batch_dims = means.shape[:-2] + B = math.prod(batch_dims) + N = means.shape[-2] + if not ctx.needs_input_grad[0]: + v_means = None + else: + if sparse_grad: + # TODO: gaussian_ids is duplicated so not ideal. + # An idea is to directly set the attribute (e.g., .sparse_grad) of + # the tensor but this requires the tensor to be leaf node only. And + # a customized optimizer would be needed in this case. + v_means = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_means, # [nnz, 3] + size=means.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[1]: + v_covars = None + else: + if sparse_grad: + v_covars = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_covars, # [nnz, 6] + size=covars.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[2]: + v_quats = None + else: + if sparse_grad: + v_quats = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_quats, # [nnz, 4] + size=quats.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[3]: + v_scales = None + else: + if sparse_grad: + v_scales = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_scales, # [nnz, 3] + size=scales.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[4]: + v_viewmats = None + + return ( + v_means, + v_covars, + v_quats, + v_scales, + v_viewmats, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +class _SphericalHarmonics(torch.autograd.Function): + """Spherical Harmonics""" + + @staticmethod + def forward( + ctx, sh_degree: int, dirs: Tensor, coeffs: Tensor, masks: Tensor + ) -> Tensor: + colors = _make_lazy_sycl_func("spherical_harmonics_fwd")( + sh_degree, dirs, coeffs, masks + ) + ctx.save_for_backward(dirs, coeffs, masks) + ctx.sh_degree = sh_degree + ctx.num_bases = coeffs.shape[-2] + return colors + + @staticmethod + def backward(ctx, v_colors: Tensor): + dirs, coeffs, masks = ctx.saved_tensors + sh_degree = ctx.sh_degree + num_bases = ctx.num_bases + compute_v_dirs = ctx.needs_input_grad[1] + v_coeffs, v_dirs = _make_lazy_sycl_func("spherical_harmonics_bwd")( + num_bases, + sh_degree, + dirs, + coeffs, + masks, + v_colors.contiguous(), + compute_v_dirs, + ) + if not compute_v_dirs: + v_dirs = None + return None, v_dirs, v_coeffs, None + + +###### 2DGS ###### +def fully_fused_projection_2dgs( + means: Tensor, # [..., N, 3] + quats: Tensor, # [..., N, 4] + scales: Tensor, # [..., N, 3] + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + eps2d: float = 0.3, + near_plane: float = 0.01, + far_plane: float = 1e10, + radius_clip: float = 0.0, + packed: bool = False, + sparse_grad: bool = False, +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Prepare Gaussians for rasterization + + This function prepares ray-splat intersection matrices, computes + per splat bounding box and 2D means in image space. + + Args: + means: Gaussian means. [..., N, 3] + quats: Quaternions (No need to be normalized). [..., N, 4]. + scales: Scales. [..., N, 3]. + viewmats: World-to-camera matrices. [..., C, 4, 4] + Ks: Camera intrinsics. [..., C, 3, 3] + width: Image width. + height: Image height. + near_plane: Near plane distance. Default: 0.01. + far_plane: Far plane distance. Default: 200. + radius_clip: Gaussians with projected radii smaller than this value will be ignored. Default: 0.0. + packed: If True, the output tensors will be packed into a flattened tensor. Default: False. + sparse_grad (Experimental): This is only effective when `packed` is True. If True, during backward the gradients + of {`means`, `covars`, `quats`, `scales`} will be a sparse Tensor in COO layout. Default: False. + + Returns: + A tuple: + + If `packed` is True: + + - **batch_ids**. The batch indices of the projected Gaussians. Int32 tensor of shape [nnz]. + - **camera_ids**. The camera indices of the projected Gaussians. Int32 tensor of shape [nnz]. + - **gaussian_ids**. The column indices of the projected Gaussians. Int32 tensor of shape [nnz]. + - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [nnz, 2]. + - **means**. Projected Gaussian means in 2D. [nnz, 2] + - **depths**. The z-depth of the projected Gaussians. [nnz] + - **ray_transforms**. transformation matrices that transforms xy-planes in pixel spaces into splat coordinates (WH)^T in equation (9) in paper [nnz, 3, 3] + - **normals**. The normals in camera spaces. [nnz, 3] + + If `packed` is False: + + - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [..., C, N, 2]. + - **means**. Projected Gaussian means in 2D. [..., C, N, 2] + - **depths**. The z-depth of the projected Gaussians. [..., C, N] + - **ray_transforms**. transformation matrices that transforms xy-planes in pixel spaces into splat coordinates [..., C, N, 3, 3] + - **normals**. The normals in camera spaces. [..., C, N, 3] + + """ + batch_dims = means.shape[:-2] + N = means.shape[-2] + C = viewmats.shape[-3] + assert means.shape == batch_dims + (N, 3), means.shape + assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape + assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape + means = means.contiguous() + assert quats is not None, "quats is required" + assert scales is not None, "scales is required" + assert quats.shape == batch_dims + (N, 4), quats.shape + assert scales.shape == batch_dims + (N, 3), scales.shape + quats = quats.contiguous() + scales = scales.contiguous() + if sparse_grad: + assert packed, "sparse_grad is only supported when packed is True" + + viewmats = viewmats.contiguous() + Ks = Ks.contiguous() + if packed: + return _FullyFusedProjectionPacked2DGS.apply( + means, + quats, + scales, + viewmats, + Ks, + width, + height, + near_plane, + far_plane, + radius_clip, + sparse_grad, + ) + else: + return _FullyFusedProjection2DGS.apply( + means, + quats, + scales, + viewmats, + Ks, + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + ) + + +class _FullyFusedProjection2DGS(torch.autograd.Function): + """Projects Gaussians to 2D.""" + + @staticmethod + def forward( + ctx, + means: Tensor, # [..., N, 3] + quats: Tensor, # [..., N, 4] + scales: Tensor, # [..., N, 3] + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + eps2d: float, + near_plane: float, + far_plane: float, + radius_clip: float, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + radii, means2d, depths, ray_transforms, normals = _make_lazy_sycl_func( + "projection_2dgs_fused_fwd" + )( + means, + quats, + scales, + viewmats, + Ks, + width, + height, + eps2d, + near_plane, + far_plane, + radius_clip, + ) + ctx.save_for_backward( + means, + quats, + scales, + viewmats, + Ks, + radii, + ray_transforms, + normals, + ) + ctx.width = width + ctx.height = height + ctx.eps2d = eps2d + + return radii, means2d, depths, ray_transforms, normals + + @staticmethod + def backward(ctx, v_radii, v_means2d, v_depths, v_ray_transforms, v_normals): + ( + means, + quats, + scales, + viewmats, + Ks, + radii, + ray_transforms, + normals, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + eps2d = ctx.eps2d + v_means, v_quats, v_scales, v_viewmats = _make_lazy_sycl_func( + "projection_2dgs_fused_bwd" + )( + means, + quats, + scales, + viewmats, + Ks, + width, + height, + radii, + ray_transforms, + v_means2d.contiguous(), + v_depths.contiguous(), + v_normals.contiguous(), + v_ray_transforms.contiguous(), + ctx.needs_input_grad[3], # viewmats_requires_grad + ) + if not ctx.needs_input_grad[0]: + v_means = None + if not ctx.needs_input_grad[1]: + v_quats = None + if not ctx.needs_input_grad[2]: + v_scales = None + if not ctx.needs_input_grad[3]: + v_viewmats = None + + return ( + v_means, + v_quats, + v_scales, + v_viewmats, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +class _FullyFusedProjectionPacked2DGS(torch.autograd.Function): + """Projects Gaussians to 2D. Return packed tensors.""" + + @staticmethod + def forward( + ctx, + means: Tensor, # [..., N, 3] + quats: Tensor, # [..., N, 4] + scales: Tensor, # [..., N, 3] + viewmats: Tensor, # [..., C, 4, 4] + Ks: Tensor, # [..., C, 3, 3] + width: int, + height: int, + near_plane: float, + far_plane: float, + radius_clip: float, + sparse_grad: bool, + ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + ( + indptr, + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + ray_transforms, + normals, + ) = _make_lazy_sycl_func("projection_2dgs_packed_fwd")( + means, + quats, + scales, + viewmats, + Ks, + width, + height, + near_plane, + far_plane, + radius_clip, + ) + ctx.save_for_backward( + batch_ids, + camera_ids, + gaussian_ids, + means, + quats, + scales, + viewmats, + Ks, + ray_transforms, + ) + ctx.width = width + ctx.height = height + ctx.sparse_grad = sparse_grad + + return ( + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + ray_transforms, + normals, + ) + + @staticmethod + def backward( + ctx, + v_batch_ids, + v_camera_ids, + v_gaussian_ids, + v_radii, + v_means2d, + v_depths, + v_ray_transforms, + v_normals, + ): + ( + batch_ids, + camera_ids, + gaussian_ids, + means, + quats, + scales, + viewmats, + Ks, + ray_transforms, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + sparse_grad = ctx.sparse_grad + + v_means, v_quats, v_scales, v_viewmats = _make_lazy_sycl_func( + "projection_2dgs_packed_bwd" + )( + means, + quats, + scales, + viewmats, + Ks, + width, + height, + batch_ids, + camera_ids, + gaussian_ids, + ray_transforms, + v_means2d.contiguous(), + v_depths.contiguous(), + v_ray_transforms.contiguous(), + v_normals.contiguous(), + ctx.needs_input_grad[3], # viewmats_requires_grad + sparse_grad, + ) + + if sparse_grad: + batch_dims = means.shape[:-2] + B = math.prod(batch_dims) + N = means.shape[-2] + + if not ctx.needs_input_grad[0]: + v_means = None + else: + if sparse_grad: + # TODO: gaussian_ids is duplicated so not ideal. + # An idea is to directly set the attribute (e.g., .sparse_grad) of + # the tensor but this requires the tensor to be leaf node only. And + # a customized optimizer would be needed in this case. + v_means = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_means, # [nnz, 3] + size=means.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[1]: + v_quats = None + else: + if sparse_grad: + v_quats = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_quats, # [nnz, 4] + size=quats.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[2]: + v_scales = None + else: + if sparse_grad: + v_scales = torch.sparse_coo_tensor( + indices=gaussian_ids[None], + values=v_scales, # [nnz, 3] + size=scales.shape, + is_coalesced=len(viewmats) == 1, + ) + if not ctx.needs_input_grad[3]: + v_viewmats = None + + return ( + v_means, + v_quats, + v_scales, + v_viewmats, + None, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def rasterize_to_pixels_2dgs( + means2d: Tensor, # [..., N, 2] + ray_transforms: Tensor, # [..., N, 3, 3] + colors: Tensor, # [..., N, channels] + opacities: Tensor, # [..., N] + normals: Tensor, # [..., N, 3] + densify: Tensor, # [..., N, 2] + image_width: int, + image_height: int, + tile_size: int, + isect_offsets: Tensor, # [..., tile_height, tile_width] + flatten_ids: Tensor, # [n_isects] + backgrounds: Optional[Tensor] = None, # [..., channels] + masks: Optional[Tensor] = None, # [..., tile_height, tile_width] + packed: bool = False, + absgrad: bool = False, + distloss: bool = False, +) -> Tuple[Tensor, Tensor]: + """Rasterize Gaussians to pixels. + + Args: + means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. + ray_transforms: transformation matrices that transforms xy-planes in pixel spaces into splat coordinates. [..., N, 3, 3] if packed is False, [nnz, channels] if packed is True. + colors: Gaussian colors or ND features. [..., N, channels] if packed is False, [nnz, channels] if packed is True. + opacities: Gaussian opacities that support per-view values. [..., N] if packed is False, [nnz] if packed is True. + normals: The normals in camera space. [..., N, 3] if packed is False, [nnz, 3] if packed is True. + densify: Dummy variable to keep track of gradient for densification. [..., N, 2] if packed, [nnz, 3] if packed is True. + tile_size: Tile size. + isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] + flatten_ids: The global flatten indices in [I * N] or [nnz] from `isect_tiles()`. [n_isects] + backgrounds: Background colors. [..., channels]. Default: None. + masks: Optional tile mask to skip rendering GS to masked tiles. [..., tile_height, tile_width]. Default: None. + packed: If True, the input tensors are expected to be packed with shape [nnz, ...]. Default: False. + absgrad: If True, the backward pass will compute a `.absgrad` attribute for `means2d`. Default: False. + + Returns: + A tuple: + + - **Rendered colors**. [..., image_height, image_width, channels] + - **Rendered alphas**. [..., image_height, image_width, 1] + - **Rendered normals**. [..., image_height, image_width, 3] + - **Rendered distortion**. [..., image_height, image_width, 1] + - **Rendered median depth**.[..., image_height, image_width, 1] + + + """ + image_dims = means2d.shape[:-2] + channels = colors.shape[-1] + device = means2d.device + if packed: + nnz = means2d.size(0) + assert means2d.shape == (nnz, 2), means2d.shape + assert ray_transforms.shape == (nnz, 3, 3), ray_transforms.shape + assert colors.shape[0] == nnz, colors.shape + assert opacities.shape == (nnz,), opacities.shape + else: + N = means2d.size(-2) + assert means2d.shape == image_dims + (N, 2), means2d.shape + assert ray_transforms.shape == image_dims + (N, 3, 3), ray_transforms.shape + assert colors.shape[:-2] == image_dims, colors.shape + assert opacities.shape == image_dims + (N,), opacities.shape + if backgrounds is not None: + assert backgrounds.shape == image_dims + (channels,), backgrounds.shape + backgrounds = backgrounds.contiguous() + + # Pad the channels to the nearest supported number if necessary + if channels > 512 or channels == 0: + # TODO: maybe worth to support zero channels? + 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 sycl kernel) + colors = torch.cat( + [ + colors[..., :-1], + torch.empty(*colors.shape[:-1], padded_channels, device=device), + colors[..., -1:], + ], + dim=-1, + ) + if backgrounds is not None: + backgrounds = torch.cat( + [ + backgrounds, + torch.zeros( + *backgrounds.shape[:-1], padded_channels, device=device + ), + ], + dim=-1, + ) + else: + padded_channels = 0 + tile_height, tile_width = isect_offsets.shape[-2:] + assert ( + tile_height * tile_size >= image_height + ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" + assert ( + tile_width * tile_size >= image_width + ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" + + ( + render_colors, + render_alphas, + render_normals, + render_distort, + render_median, + ) = _RasterizeToPixels2DGS.apply( + means2d.contiguous(), + ray_transforms.contiguous(), + colors.contiguous(), + opacities.contiguous(), + normals.contiguous(), + densify.contiguous(), + backgrounds, + masks, + image_width, + image_height, + tile_size, + isect_offsets.contiguous(), + flatten_ids.contiguous(), + absgrad, + distloss, + ) + + if padded_channels > 0: + render_colors = torch.cat( + [render_colors[..., : -padded_channels - 1], render_colors[..., -1:]], + dim=-1, + ) + + return render_colors, render_alphas, render_normals, render_distort, render_median + + +@torch.no_grad() +def rasterize_to_indices_in_range_2dgs( + range_start: int, + range_end: int, + transmittances: Tensor, # [..., image_height, image_width] + means2d: Tensor, # [..., N, 2] + ray_transforms: Tensor, # [..., N, 3, 3] + opacities: Tensor, # [..., N] + image_width: int, + image_height: int, + tile_size: int, + isect_offsets: Tensor, + flatten_ids: Tensor, +) -> Tuple[Tensor, Tensor, Tensor]: + """Rasterizes a batch of Gaussians to images but only returns the indices. + + .. note:: + + This function supports iterative rasterization, in which each call of this function + will rasterize a batch of Gaussians from near to far, defined by `[range_start, range_end)`. + If a one-step full rasterization is desired, set `range_start` to 0 and `range_end` to a really + large number, e.g, 1e10. + + Args: + range_start: The start batch of Gaussians to be rasterized (inclusive). + range_end: The end batch of Gaussians to be rasterized (exclusive). + transmittances: Currently transmittances. [..., image_height, image_width] + means2d: Projected Gaussian means. [..., N, 2] + ray_transforms: transformation matrices that transforms xy-planes in pixel spaces into splat coordinates. [..., N, 3, 3] + opacities: Gaussian opacities that support per-view values. [..., N] + image_width: Image width. + image_height: Image height. + tile_size: Tile size. + isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] + flatten_ids: The global flatten indices in [I * N] from `isect_tiles()`. [n_isects] + + Returns: + A tuple: + + - **Gaussian ids**. Gaussian ids for the pixel intersection. A flattened list of shape [M]. + - **Pixel ids**. pixel indices (row-major). A flattened list of shape [M]. + - **Camera ids**. Camera indices. A flattened list of shape [M]. + - **Batch ids**. Batch indices. A flattened list of shape [M]. + """ + + image_dims = means2d.shape[:-2] + tile_height, tile_width = isect_offsets.shape[-2:] + N = means2d.shape[-2] + assert transmittances.shape == image_dims + ( + image_height, + image_width, + ), transmittances.shape + assert means2d.shape == image_dims + (N, 2), means2d.shape + assert ray_transforms.shape == image_dims + (N, 3, 3), ray_transforms.shape + assert opacities.shape == image_dims + (N,), opacities.shape + assert isect_offsets.shape == image_dims + ( + tile_height, + tile_width, + ), isect_offsets.shape + assert ( + tile_height * tile_size >= image_height + ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" + assert ( + tile_width * tile_size >= image_width + ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" + + out_gauss_ids, out_indices = _make_lazy_sycl_func("rasterize_to_indices_2dgs")( + range_start, + range_end, + transmittances.contiguous(), + means2d.contiguous(), + ray_transforms.contiguous(), + opacities.contiguous(), + image_width, + image_height, + tile_size, + isect_offsets.contiguous(), + flatten_ids.contiguous(), + ) + out_pixel_ids = out_indices % (image_width * image_height) + out_image_ids = out_indices // (image_width * image_height) + return out_gauss_ids, out_pixel_ids, out_image_ids + + +class _RasterizeToPixels2DGS(torch.autograd.Function): + """Rasterize gaussians 2DGS""" + + @staticmethod + def forward( + ctx, + means2d: Tensor, + ray_transforms: Tensor, + colors: Tensor, + opacities: Tensor, + normals: Tensor, + densify: Tensor, + backgrounds: Tensor, + masks: Tensor, + width: int, + height: int, + tile_size: int, + isect_offsets: Tensor, + flatten_ids: Tensor, + absgrad: bool, + distloss: bool, + ) -> Tuple[Tensor, Tensor]: + ( + render_colors, + render_alphas, + render_normals, + render_distort, + render_median, + last_ids, + median_ids, + ) = _make_lazy_sycl_func("rasterize_to_pixels_2dgs_fwd")( + means2d, + ray_transforms, + colors, + opacities, + normals, + backgrounds, + masks, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + ) + + ctx.save_for_backward( + means2d, + ray_transforms, + colors, + opacities, + normals, + densify, + backgrounds, + masks, + isect_offsets, + flatten_ids, + render_colors, + render_alphas, + last_ids, + median_ids, + ) + ctx.width = width + ctx.height = height + ctx.tile_size = tile_size + ctx.absgrad = absgrad + ctx.distloss = distloss + + # double to float + render_alphas = render_alphas.float() + return ( + render_colors, + render_alphas, + render_normals, + render_distort, + render_median, + ) + + @staticmethod + def backward( + ctx, + v_render_colors: Tensor, + v_render_alphas: Tensor, + v_render_normals: Tensor, + v_render_distort: Tensor, + v_render_median: Tensor, + ): + + ( + means2d, + ray_transforms, + colors, + opacities, + normals, + densify, + backgrounds, + masks, + isect_offsets, + flatten_ids, + render_colors, + render_alphas, + last_ids, + median_ids, + ) = ctx.saved_tensors + width = ctx.width + height = ctx.height + tile_size = ctx.tile_size + absgrad = ctx.absgrad + + ( + v_means2d_abs, + v_means2d, + v_ray_transforms, + v_colors, + v_opacities, + v_normals, + v_densify, + ) = _make_lazy_sycl_func("rasterize_to_pixels_2dgs_bwd")( + means2d, + ray_transforms, + colors, + opacities, + normals, + densify, + backgrounds, + masks, + width, + height, + tile_size, + isect_offsets, + flatten_ids, + render_colors, + render_alphas, + last_ids, + median_ids, + v_render_colors.contiguous(), + v_render_alphas.contiguous(), + v_render_normals.contiguous(), + v_render_distort.contiguous(), + v_render_median.contiguous(), + absgrad, + ) + torch.sycl.synchronize() + if absgrad: + means2d.absgrad = v_means2d_abs + + if ctx.needs_input_grad[6]: + v_backgrounds = (v_render_colors * (1.0 - render_alphas).float()).sum( + dim=(-3, -2) + ) + else: + v_backgrounds = None + + return ( + v_means2d, + v_ray_transforms, + v_colors, + v_opacities, + v_normals, + v_densify, + v_backgrounds, + None, + None, + None, + None, + None, + None, + None, + None, + ) diff --git a/gsplat/sycl/ext.cpp b/gsplat/sycl/ext.cpp new file mode 100644 index 00000000..144a889f --- /dev/null +++ b/gsplat/sycl/ext.cpp @@ -0,0 +1,104 @@ +#include + +#include "Ops.h" +#include "Cameras.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + + py::enum_< gsplat::xpu::CameraModelType>(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..f99a937c --- /dev/null +++ b/gsplat/sycl/include/Cameras.h @@ -0,0 +1,58 @@ +#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..df56d625 --- /dev/null +++ b/gsplat/sycl/include/Common.h @@ -0,0 +1,32 @@ +#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_CONTIGUOUS(x) \ + TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") +#define CHECK_INPUT(x) \ + CHECK_XPU(x); \ + CHECK_CONTIGUOUS(x) + + +// +// Legacy Camera Types +// +enum CameraModelType { + PINHOLE = 0, + ORTHO = 1, + FISHEYE = 2, + FTHETA = 3, +}; + +#define GSPLAT_N_THREADS 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..bb8bb37c --- /dev/null +++ b/gsplat/sycl/include/Ops.h @@ -0,0 +1,567 @@ +// A collection of operators for gsplat +#pragma once + +#include +#include +#include "Cameras.h" +#include "Common.h" +#include "types.hpp" + +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/gsplat_sycl_utils.hpp b/gsplat/sycl/include/gsplat_sycl_utils.hpp new file mode 100644 index 00000000..76bf074f --- /dev/null +++ b/gsplat/sycl/include/gsplat_sycl_utils.hpp @@ -0,0 +1,75 @@ +#ifndef GSPLAT_SYCL_UTILS +#define GSPLAT_SYCL_UTILS + + +#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< const T *>(source)); +} + +template +void gpuAtomicAdd(T* ptr, T value) { + sycl::atomic_ref + protected_ref(*ptr); + protected_ref.fetch_add(value); +} + +template +void gpuAtomicAddGlobal(T& ref, const T& value) { + sycl::atomic_ref + protected_ref(ref); + protected_ref.fetch_add(value); +} + +template +void gpuAtomicAddLocal(T& ref, const T& value) { + sycl::atomic_ref + protected_ref(ref); + protected_ref.fetch_add(value); +} + +#endif \ No newline at end of file diff --git a/gsplat/sycl/include/helpers.hpp b/gsplat/sycl/include/helpers.hpp new file mode 100644 index 00000000..4372d76d --- /dev/null +++ b/gsplat/sycl/include/helpers.hpp @@ -0,0 +1,14 @@ +#ifndef GSPLAT_SYCL_HELPERS_HPP +#define GSPLAT_SYCL_HELPERS_HPP + +#include + +template +void gpuAtomicAdd(T* ptr, T value) { + sycl::atomic_ref + protected_ref(*ptr); + protected_ref.fetch_add(value); +} + +#endif //GSPLAT_SYCL_HELPERS_HPP \ 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..1859e414 --- /dev/null +++ b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp @@ -0,0 +1,69 @@ +#ifndef ComputeShBwdKernel_HPP +#define ComputeShBwdKernel_HPP + +#include "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); + } + } +}; + +#endif //ComputeShBwdKernel_HPP + +} // 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..963326ba --- /dev/null +++ b/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp @@ -0,0 +1,55 @@ +#ifndef ComputeShFwdKernel_HPP +#define ComputeShFwdKernel_HPP + +#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 + ); + } +}; + +#endif //ComputeShFwdKernel_HPP + +} // 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..b3f72703 --- /dev/null +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp @@ -0,0 +1,273 @@ +#ifndef FullyFusedProjectionBwdKernel_HPP +#define FullyFusedProjectionBwdKernel_HPP + + +#include "utils.hpp" +#include "quat.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "proj.hpp" +#include "transform.hpp" + +namespace gsplat::xpu { + +template +struct FullyFusedProjectionBwdKernel{ + // fwd inputs + const uint32_t m_C; + const uint32_t m_N; + const T* m_means; // [N, 3] + const T* m_covars; // [N, 6] optional + const T* m_quats; // [N, 4] optional + const T* m_scales; // [N, 3] optional + const T* m_viewmats; // [C, 4, 4] + const T* m_Ks; // [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 + const int32_t* m_radii; // [C, N] + const T* m_conics; // [C, N, 3] + const T* m_compensations; // [C, N] optional + // grad outputs + const T* m_v_means2d; // [C, N, 2] + const T* m_v_depths; // [C, N] + const T* m_v_conics; // [C, N, 3] + const T* m_v_compensations; // [C, N] optional + // grad inputs + T* m_v_means; // [N, 3] + T* m_v_covars; // [N, 6] optional + T* m_v_quats; // [N, 4] optional + T* m_v_scales; // [N, 3] optional + T* m_v_viewmats;// [C, 4, 4] optional + + FullyFusedProjectionBwdKernel( + 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 + ) + : 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); + if (idx >= m_C * m_N || m_radii[idx] <= 0) { + 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* viewmats = m_viewmats + (cid * 16); + const T* Ks = m_Ks + (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) { + const T* covars = m_covars + (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 + quat = glm::make_vec4(m_quats + (gid * 4)); + scale = glm::make_vec3(m_scales + (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) { + 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 * 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 + ); + T* v_quats = m_v_quats + (gid * 4); + T* v_scales = m_v_scales + (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) { + 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]); + } + } + } +}; + +#endif //FullyFusedProjectionBwdKernel_HPP + +} // namespace gsplat::xpu diff --git a/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp new file mode 100644 index 00000000..d248d14b --- /dev/null +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp @@ -0,0 +1,220 @@ +#ifndef FullyFusedProjectionFwdKernel_HPP +#define FullyFusedProjectionFwdKernel_HPP + + +#include "utils.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "proj.hpp" +#include "transform.hpp" + +namespace gsplat::xpu { + +template +struct FullyFusedProjectionFwdKernel{ + const uint32_t m_C; + const uint32_t m_N; + const T* m_means; // [N, 3] + const T* m_covars; // [N, 6] optional + const T* m_quats; // [N, 4] optional + const T* m_scales; // [N, 3] optional + const T* m_viewmats; // [C, 4, 4] + const T* m_Ks; // [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 + int32_t * m_radii; // [C, N] + T* m_means2d; // [C, N, 2] + T* m_depths; // [C, N] + T* m_conics; // [C, N, 3] + T* m_compensations; // [C, N] optional + + FullyFusedProjectionFwdKernel( + 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 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 + ) + : 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_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); + 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 + + const T* means = m_means + (gid * 3); + const T* viewmats = m_viewmats + (cid * 16); + const T* Ks = m_Ks + (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) { + m_radii[idx] = 0; + return; + } + + // transform Gaussian covariance to camera space + mat3 covar; + if (m_covars != nullptr) { + const T* covars = m_covars + (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 + const T* quats = m_quats + (gid * 4); + const T* scales = m_scales + (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) { + m_radii[idx] = 0; + return; + } + + // compute the inverse of the 2d covariance + mat2 covar2d_inv; + inverse(covar2d, covar2d_inv); + + // take 3 sigma as the radius (non differentiable) + T b = 0.5f * (covar2d[0][0] + covar2d[1][1]); + T v1 = b + sycl::sqrt(sycl::max(0.01f, b * b - det)); + T radius = sycl::ceil(3.f * sycl::sqrt(v1)); + + if (radius <= m_radius_clip) { + m_radii[idx] = 0; + return; + } + + // mask out gaussians outside the image region + if (mean2d.x + radius <= 0 || mean2d.x - radius >= m_image_width || + mean2d.y + radius <= 0 || mean2d.y - radius >= m_image_height) { + m_radii[idx] = 0; + return; + } + + // write to outputs + m_radii[idx] = (int32_t)radius; + 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; + } + + } + +}; +#endif //FullyFusedProjectionFwdKernel_HPP + +} // 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..a6fba80c --- /dev/null +++ b/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp @@ -0,0 +1,72 @@ +#ifndef IsectOffsetEncodeKernel_HPP +#define IsectOffsetEncodeKernel_HPP + +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); + } + } +}; + +#endif + +} // 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..ba6bb81a --- /dev/null +++ b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp @@ -0,0 +1,143 @@ +#ifndef IsectTilesKernel_HPP +#define IsectTilesKernel_HPP + + +#include "types.hpp" +#include "transform.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 = m_radii[idx]; + if (radius <= 0) { + if (first_pass) { + m_tiles_per_gauss[idx] = 0; + } + return; + } + + vec2 mean2d = glm::make_vec2(m_means2d + 2 * idx); + + T tile_radius = radius / 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; + tile_min.x = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::floor(tile_x - tile_radius)), m_tile_width); + tile_min.y = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::floor(tile_y - tile_radius)), m_tile_height); + + tile_max.x = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::ceil(tile_x + tile_radius)), m_tile_width); + tile_max.y = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::ceil(tile_y + tile_radius)), 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); + + int64_t depth_id_enc = (int64_t) * (int32_t *)&(m_depths[idx]); + 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; + } + } + } +}; + +#endif //IsectTilesKernel_HPP + +} // 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..ffe84fe0 --- /dev/null +++ b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp @@ -0,0 +1,140 @@ +#ifndef ProjBwdKernel_HPP +#define ProjBwdKernel_HPP + + +#include "proj.hpp" +#include "Common.h" + +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); + if (idx >= m_C * m_N) { + return; + } + + const uint32_t cid = idx / m_N; // camera id + + // shift pointers to the current camera and gaussian + 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); + const T* Ks = m_Ks + (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]); + } + } +}; + +#endif //ProjBwdKernel_HPP + +} // 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..c50b6229 --- /dev/null +++ b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp @@ -0,0 +1,88 @@ +#ifndef ProjFwdKernel_HPP +#define ProjFwdKernel_HPP + + +#include "proj.hpp" +#include "Common.h" + +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); + if (idx >= m_C * m_N) { + return; + } + const uint32_t cid = idx / m_N; // camera id + + const T* means = m_means + (idx * 3); + const T* covars = m_covars + (idx * 9); + const T* Ks = m_Ks + (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]); + } + } + +}; +#endif //ProjFwdKernel_HPP + +} //namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp new file mode 100644 index 00000000..1baf9400 --- /dev/null +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp @@ -0,0 +1,121 @@ +#ifndef QuatScaleToCovarPreciBwdKernel_HPP +#define QuatScaleToCovarPreciBwdKernel_HPP + + +#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]); + } + } + +}; + +#endif //QuatScaleToCovarPreciBwdKernel_HPP + +} // 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..a82961f1 --- /dev/null +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp @@ -0,0 +1,94 @@ +#ifndef QuatScaleToCovarPreciFwdKernel_HPP +#define QuatScaleToCovarPreciFwdKernel_HPP + + +#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]); + } + } + } + } + } +}; +#endif //QuatScaleToCovarPreciFwdKernel_HPP + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp new file mode 100644 index 00000000..75c237ee --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp @@ -0,0 +1,370 @@ +#ifndef RasterizeToPixelsBwdKernel_HPP +#define RasterizeToPixelsBwdKernel_HPP + + +#include "types.hpp" +#include "gsplat_sycl_utils.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(); + 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(); + + auto conic = *(reinterpret_cast*>(data+2) ); + m_slm_conics[threadRank] = conic.template convert(); + + if constexpr(BufferType::isVec && COLOR_DIM <= 4){ + auto color = *( reinterpret_cast*>(data + 2 + 3) ); + m_slm_colors[threadRank] = color.template convert();; + } + } + m_slm_opacities[threadRank] = static_cast(*(data + 2 + 3 + COLOR_DIM)); + + } else { + m_slm_means2d[threadRank] = m_means2d[g].template convert(); + + m_slm_opacities[threadRank] = static_cast(m_opacities[g]); + auto temp = *( reinterpret_cast*>(m_conics + g) ); + + m_slm_conics[threadRank] = temp.template convert(); + + if constexpr(BufferType::isVec && COLOR_DIM <= 4){ + auto temp2 = *( reinterpret_cast*>(m_colors + g * COLOR_DIM) ); + m_slm_colors[threadRank] = temp2.template convert(); + } + } + } + + 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(); + } 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]); + } + } + } + } + } +}; + +#endif // RasterizeToPixelsBwdKernel_HPP + +} // 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..a25718a8 --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp @@ -0,0 +1,269 @@ +#ifndef RasterizeToPixelsFwdKernel_HPP +#define RasterizeToPixelsFwdKernel_HPP + +#include "types.hpp" +#include "gsplat_sycl_utils.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*>(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]; + } + } + } +}; + +#endif //RasterizeToPixelsFwdKernel_HPP + +} // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp new file mode 100644 index 00000000..004af61f --- /dev/null +++ b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp @@ -0,0 +1,124 @@ +#ifndef WorldToCamBwdKernel_HPP +#define WorldToCamBwdKernel_HPP + + +#include "types.hpp" +#include "transform.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]); + } + } + } +}; + +#endif //WorldToCamBwdKernel_HPP + +} // 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..9296a715 --- /dev/null +++ b/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp @@ -0,0 +1,98 @@ +#ifndef WorldToCamFwdKernel_HPP +#define WorldToCamFwdKernel_HPP + +/**************************************************************************** + * World to Camera Transformation Forward Pass + * From: https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/csrc/world_to_cam_fwd.cu + ****************************************************************************/ + +#include "types.hpp" +#include "transform.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]); + } + } + } + } +}; + +#endif //WorldToCamFwdKernel_HPP + +} // 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..a5e5505c --- /dev/null +++ b/gsplat/sycl/include/proj.hpp @@ -0,0 +1,345 @@ +#ifndef GSPLAT_SYCL_PROJ_HPP +#define GSPLAT_SYCL_PROJ_HPP + +#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; +} + + +#endif // GSPLAT_SYCL_PROJ_HPP diff --git a/gsplat/sycl/include/quat.hpp b/gsplat/sycl/include/quat.hpp new file mode 100644 index 00000000..049b8efc --- /dev/null +++ b/gsplat/sycl/include/quat.hpp @@ -0,0 +1,59 @@ +#ifndef GSPLAT_SYCL_QUAT_HPP +#define GSPLAT_SYCL_QUAT_HPP + +#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; +} + + +#endif // GSPLAT_SYCL_QUAT_HPP 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..48dc2e07 --- /dev/null +++ b/gsplat/sycl/include/quat_scale_to_covar_preci.hpp @@ -0,0 +1,125 @@ +#ifndef GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP +#define GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP + +#include "types.hpp" +#include "quat.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]); +} + +#endif // GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP diff --git a/gsplat/sycl/include/spherical_harmonics.hpp b/gsplat/sycl/include/spherical_harmonics.hpp new file mode 100644 index 00000000..45bcfdbb --- /dev/null +++ b/gsplat/sycl/include/spherical_harmonics.hpp @@ -0,0 +1,362 @@ +#ifndef GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP +#define GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP + +#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; + } +} + + +#endif // GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP \ 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..e8e281eb --- /dev/null +++ b/gsplat/sycl/include/transform.hpp @@ -0,0 +1,69 @@ +#ifndef GSPLAT_SYCL_TRANSFORM_HPP +#define GSPLAT_SYCL_TRANSFORM_HPP + +#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; +} + +#endif // GSPLAT_SYCL_TRANSFORM_HPP diff --git a/gsplat/sycl/include/types.hpp b/gsplat/sycl/include/types.hpp new file mode 100644 index 00000000..8048bd83 --- /dev/null +++ b/gsplat/sycl/include/types.hpp @@ -0,0 +1,22 @@ +#ifndef GSPLAT_SYCL_TYPES_HPP +#define GSPLAT_SYCL_TYPES_HPP + +#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>; + + +#endif // GSPLAT_SYCL_TYPES_HPP \ 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..6d63b3ac --- /dev/null +++ b/gsplat/sycl/include/utils.hpp @@ -0,0 +1,84 @@ +#ifndef GSPLAT_SYCL_UTILS_HPP +#define GSPLAT_SYCL_UTILS_HPP + +#include "types.hpp" + +#include + +template +void gpuAtomicAdd(T* ptr, T value) { + sycl::atomic_ref + protected_ref(*ptr); + protected_ref.fetch_add(value); +} + +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(0.f, 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); +} + + +#endif // GSPLAT_SYCL_UTILS_HPP diff --git a/gsplat/sycl/src/adam.cpp b/gsplat/sycl/src/adam.cpp new file mode 100644 index 00000000..fd030547 --- /dev/null +++ b/gsplat/sycl/src/adam.cpp @@ -0,0 +1,23 @@ + +#include + +#include "Ops.h" +#include "Common.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..622001f0 --- /dev/null +++ b/gsplat/sycl/src/intersect_offset.cpp @@ -0,0 +1,58 @@ +#include + +#include + +#include "Ops.h" +#include "Common.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 +) { + CHECK_CONTIGUOUS(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..1fda9580 --- /dev/null +++ b/gsplat/sycl/src/intersect_tile.cpp @@ -0,0 +1,121 @@ +#include +#include + +#include "Ops.h" +#include "Common.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 +) { + CHECK_CONTIGUOUS(means2d); + CHECK_CONTIGUOUS(radii); + CHECK_CONTIGUOUS(depths); + if (image_ids.has_value()) CHECK_CONTIGUOUS(image_ids.value()); + if (gaussian_ids.has_value()) CHECK_CONTIGUOUS(gaussian_ids.value()); + + 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(1); + 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) + ); + } + + at::Tensor tiles_per_gauss = at::empty_like(depths, 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}, at::kLong); + at::Tensor flatten_ids = at::empty({n_isects}, 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..88310566 --- /dev/null +++ b/gsplat/sycl/src/null.cpp @@ -0,0 +1,13 @@ + +#include + +#include "Ops.h" +#include "Common.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..945ad018 --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp @@ -0,0 +1,32 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +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 +) { + 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_fwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp new file mode 100644 index 00000000..c5167a00 --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp @@ -0,0 +1,31 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +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 +) { + 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_bwd.cpp b/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp new file mode 100644 index 00000000..cb20c2a3 --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp @@ -0,0 +1,35 @@ + +#include + +#include "Ops.h" +#include "Common.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..98cc43f9 --- /dev/null +++ b/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp @@ -0,0 +1,34 @@ + +#include + +#include "Ops.h" +#include "Common.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..be768bf6 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp @@ -0,0 +1,36 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +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 +) { + 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_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp new file mode 100644 index 00000000..c1eb4d05 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp @@ -0,0 +1,35 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +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 +) { + 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_packed_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp new file mode 100644 index 00000000..c2b6606c --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp @@ -0,0 +1,39 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +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 +) { + 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_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp new file mode 100644 index 00000000..855e85e9 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -0,0 +1,39 @@ + +#include + +#include "Ops.h" +#include "Common.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_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 +) { + 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_simple_bwd.cpp b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp new file mode 100644 index 00000000..a3a5e4b2 --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp @@ -0,0 +1,70 @@ +#include + +#include "Ops.h" +#include "Common.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] +) { + + CHECK_CONTIGUOUS(means); + CHECK_CONTIGUOUS(covars); + CHECK_CONTIGUOUS(Ks); + CHECK_CONTIGUOUS(v_means2d); + CHECK_CONTIGUOUS(v_covars2d); + + + const uint32_t C = means.size(-3); + const uint32_t N = means.size(-2); + + + at::Tensor v_means = at::empty({C, N, 3}, means.options()); + at::Tensor v_covars = at::empty({C, N, 3, 3}, covars.options()); + + + if (C > 0 && N > 0) { + auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + + + size_t numWorkGrps = (C * 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) + { + 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..7179babe --- /dev/null +++ b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp @@ -0,0 +1,63 @@ +#include + +#include "Ops.h" +#include "Common.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 +) { + CHECK_CONTIGUOUS(means); + CHECK_CONTIGUOUS(covars); + CHECK_CONTIGUOUS(Ks); + 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); + + at::Tensor means2d = at::empty({C, N, 2}, means.options()); + at::Tensor covars2d = at::empty({C, N, 2, 2}, covars.options()); + + if (C > 0 && N > 0) { + auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + + + size_t numWorkGrps = (C * 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) + { + 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..254a6f42 --- /dev/null +++ b/gsplat/sycl/src/projection_ut_3dgs_fused.cpp @@ -0,0 +1,43 @@ + +#include + +#include "Ops.h" +#include "Common.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..a023b498 --- /dev/null +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -0,0 +1,60 @@ +#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] +) { + CHECK_CONTIGUOUS(quats); + CHECK_CONTIGUOUS(scales); + if (v_covars.has_value()) { + CHECK_CONTIGUOUS(v_covars.value()); + } + if (v_precis.has_value()) { + CHECK_CONTIGUOUS(v_precis.value()); + } + 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.size(0); + 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); + + 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); + } + ); + + 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..9eb92045 --- /dev/null +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -0,0 +1,65 @@ +#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 +) { + CHECK_CONTIGUOUS(quats); + CHECK_CONTIGUOUS(scales); + TORCH_CHECK(compute_covar || compute_preci, "Must compute at least one of covar or preci"); + + const int64_t N = quats.size(0); + auto options = quats.options(); + + at::Tensor covars; + if (compute_covar) { + covars = triu ? at::empty({N, 6}, options) : at::empty({N, 3, 3}, options); + } else { + covars = at::empty({0}, options); + } + + at::Tensor precis; + if (compute_preci) { + precis = triu ? at::empty({N, 6}, options) : at::empty({N, 3, 3}, 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); + + 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); + } + ); + + 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..e1c6f224 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp @@ -0,0 +1,28 @@ + +#include + +#include "Ops.h" +#include "Common.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..a9594011 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp @@ -0,0 +1,28 @@ + +#include + +#include "Ops.h" +#include "Common.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..2ad95e89 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -0,0 +1,51 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +namespace gsplat::xpu { + +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 +) { + 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_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp new file mode 100644 index 00000000..4a479d11 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp @@ -0,0 +1,37 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +namespace gsplat::xpu { + +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] +) { + 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_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp new file mode 100644 index 00000000..de030b6d --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp @@ -0,0 +1,195 @@ +#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, + 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 + bool absgrad +) { + CHECK_CONTIGUOUS(means2d); + CHECK_CONTIGUOUS(conics); + CHECK_CONTIGUOUS(colors); + CHECK_CONTIGUOUS(opacities); + CHECK_CONTIGUOUS(tile_offsets); + CHECK_CONTIGUOUS(flatten_ids); + CHECK_CONTIGUOUS(render_alphas); + CHECK_CONTIGUOUS(last_ids); + CHECK_CONTIGUOUS(v_render_colors); + CHECK_CONTIGUOUS(v_render_alphas); + if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + + // --- 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(1); + 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..953b0e72 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp @@ -0,0 +1,154 @@ +#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, + 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 +) { + CHECK_CONTIGUOUS(means2d); + CHECK_CONTIGUOUS(conics); + CHECK_CONTIGUOUS(colors); + CHECK_CONTIGUOUS(opacities); + CHECK_CONTIGUOUS(tile_offsets); + CHECK_CONTIGUOUS(flatten_ids); + if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + + 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(1); + 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::Tensor renders = at::empty({C, image_height, image_width, channels}, options_float); + at::Tensor alphas = at::empty({C, image_height, image_width, 1}, options_float); + at::Tensor last_ids = at::empty({C, image_height, image_width}, 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..bc3faba3 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp @@ -0,0 +1,49 @@ + +#include + +#include "Ops.h" +#include "Common.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..4c3d0758 --- /dev/null +++ b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp @@ -0,0 +1,43 @@ + +#include + +#include "Ops.h" +#include "Common.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..f51efa16 --- /dev/null +++ b/gsplat/sycl/src/relocation.cpp @@ -0,0 +1,19 @@ + +#include + +#include "Ops.h" +#include "Common.h" + +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 +) { + 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/spherical_harmonics_bwd.cpp b/gsplat/sycl/src/spherical_harmonics_bwd.cpp new file mode 100644 index 00000000..bc220c16 --- /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 +) { + CHECK_CONTIGUOUS(dirs); + CHECK_CONTIGUOUS(coeffs); + CHECK_CONTIGUOUS(v_colors); + if (masks.has_value()) { + CHECK_CONTIGUOUS(masks.value()); + } + + 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); + + 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); + } + ); + + 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..83c66f9d --- /dev/null +++ b/gsplat/sycl/src/spherical_harmonics_fwd.cpp @@ -0,0 +1,59 @@ +#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 // [...] +) { + TORCH_CHECK(dirs.is_contiguous(), "Input 'dirs' tensor must be contiguous."); + TORCH_CHECK(coeffs.is_contiguous(), "Input 'coeffs' tensor must be contiguous."); + if (masks.has_value()) { + TORCH_CHECK(masks.value().is_contiguous(), "Input 'masks' tensor must be contiguous."); + } + + 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/setup.py b/setup.py index c0ba3c48..c7fd6737 100644 --- a/setup.py +++ b/setup.py @@ -6,12 +6,36 @@ import sys from setuptools import find_packages, setup +import subprocess __version__ = None exec(open("gsplat/version.py", "r").read()) URL = "https://github.com/nerfstudio-project/gsplat" +has_cuda = False +try: + import torch + has_cuda = torch.cuda.is_available() +except ImportError: + pass + +has_xpu = False +if not has_cuda: + try: + import torch + has_xpu = torch.xpu.is_available() + except (ImportError, AttributeError): + pass + +has_sycl_compiler = False +if os.system('icpx --version > /dev/null 2>&1') == 0: + has_sycl_compiler = True +elif os.system('dpcpp --version > /dev/null 2>&1') == 0: + has_sycl_compiler = True + +BUILD_SYCL = has_xpu and has_sycl_compiler + BUILD_NO_CUDA = os.getenv("BUILD_NO_CUDA", "0") == "1" WITH_SYMBOLS = os.getenv("WITH_SYMBOLS", "0") == "1" LINE_INFO = os.getenv("LINE_INFO", "0") == "1" @@ -23,6 +47,30 @@ print(f"Setting MAX_JOBS to {os.environ['MAX_JOBS']}") +from torch.utils.cpp_extension import BuildExtension + +class SyclBuildExtension(BuildExtension): + """ + Custom build class to orchestrate a CMake build for the SYCL backend. + """ + def run(self): + print("--- Running SYCL build via CMake ---") + 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) + + subprocess.check_call( + ["cmake", f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={os.path.join(install_dir, 'gsplat')}", sycl_dir], + cwd=build_dir + ) + subprocess.check_call( + ["cmake", "--build", ".", "--config", "Release", "--", f"-j{jobs}"], + cwd=build_dir + ) + def get_ext(): from torch.utils.cpp_extension import BuildExtension @@ -105,14 +153,29 @@ 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": SyclBuildExtension} + ext_modules.append(Extension("gsplat.gsplat_sycl_kernels", sources=[])) +elif not BUILD_NO_CUDA: + print("--- Configuring for CUDA build ---") + ext_modules = get_extensions() + cmdclass = {"build_ext": get_ext()} +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", # Updated to match your CMake install_requires=[ "ninja", "numpy", @@ -122,7 +185,6 @@ def get_extensions(): "typing_extensions; python_version<'3.8'", ], extras_require={ - # dev dependencies. Install them by `pip install gsplat[dev]` "dev": [ "black[jupyter]==22.3.0", "isort==5.10.1", @@ -134,12 +196,14 @@ def get_extensions(): "build", "twine", ], + "sycl": ["pybind11>=2.10"], }, - 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_basic.py b/tests/test_basic.py index fd15933a..a70db2b5 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -1,9 +1,10 @@ -"""Tests for the functions in the CUDA extension. +"""Tests for the backend functions. Usage: -```bash + pytest -s -``` + +# To force a specific backend for testing: """ import math @@ -13,15 +14,23 @@ 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 = torch.device("cpu") + +requires_backend = pytest.mark.skipif(gsplat.BACKEND == "", reason="No CUDA or SYCL 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 +43,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,218 +59,126 @@ 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 + 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), - ) - _v_quats, _v_scales = torch.autograd.grad( - (_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)) + _v_quats, _v_scales = torch.autograd.grad((_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 ( - _fisheye_proj, - _ortho_proj, - _persp_proj, - _world_to_cam, - ) - from gsplat.cuda._wrapper import proj, quat_scale_to_covar_preci +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) 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) - if camera_model == "ortho": - _means2d, _covars2d = _ortho_proj(means, covars, Ks, width, height) - elif camera_model == "fisheye": - _means2d, _covars2d = _fisheye_proj(means, covars, Ks, width, height) - elif camera_model == "pinhole": - _means2d, _covars2d = _persp_proj(means, covars, Ks, width, height) - else: - assert_never(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": _means2d, _covars2d = _fisheye_proj(means, covars, Ks, width, height) + elif camera_model == "pinhole": _means2d, _covars2d = _persp_proj(means, covars, Ks, width, height) + else: assert_never(camera_model) 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_means, v_covars = torch.autograd.grad( - (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), - ) + 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)) + _v_means, _v_covars = torch.autograd.grad((_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]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) -def test_projection( - test_data, - fused: bool, - calc_compensations: bool, - camera_model: Literal["pinhole", "ortho", "fisheye"], - 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 +def test_projection(test_data, fused: bool, calc_compensations: bool, camera_model: str, batch_dims: Tuple[int, ...]): - 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( - means, - None, - quats, - scales, - viewmats, - Ks, - width, - height, - calc_compensations=calc_compensations, - camera_model=camera_model, + radii, means2d, depths, conics, compensations = gsplat.fully_fused_projection( + means, None, quats, scales, viewmats, Ks, width, height, + calc_compensations=calc_compensations, 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( - means, - covars, - None, - None, - viewmats, - Ks, - width, - height, - calc_compensations=calc_compensations, - camera_model=camera_model, + covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=True) + radii, means2d, depths, conics, compensations = gsplat.fully_fused_projection( + means, covars, None, None, viewmats, Ks, width, height, + 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, - viewmats, - Ks, - width, - height, - calc_compensations=calc_compensations, - camera_model=camera_model, + means, _covars, viewmats, Ks, width, height, + calc_compensations=calc_compensations, 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) torch.testing.assert_close(depths[valid], _depths[valid], rtol=1e-4, atol=1e-4) torch.testing.assert_close(conics[valid], _conics[valid], rtol=1e-4, atol=1e-4) if calc_compensations: - torch.testing.assert_close( - compensations[valid], _compensations[valid], rtol=1e-4, atol=1e-3 - ) + torch.testing.assert_close(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( - (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), - ) - _v_viewmats, _v_quats, _v_scales, _v_means = torch.autograd.grad( - (_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), - ) + 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) + 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) + _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) torch.testing.assert_close(v_quats, _v_quats, rtol=2e-1, atol=2e-2) @@ -268,175 +186,75 @@ 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]) @pytest.mark.parametrize("camera_model", ["pinhole", "ortho", "fisheye"]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) -def test_fully_fused_projection_packed( - test_data, - fused: bool, - sparse_grad: bool, - calc_compensations: bool, - camera_model: Literal["pinhole", "ortho", "fisheye"], - batch_dims: Tuple[int, ...], -): - from gsplat.cuda._wrapper import fully_fused_projection, quat_scale_to_covar_preci - +def test_fully_fused_projection_packed(test_data, fused: bool, sparse_grad: bool, calc_compensations: bool, camera_model: str, batch_dims: Tuple[int, ...]): + 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, - camera_ids, - gaussian_ids, - radii, - means2d, - depths, - conics, - compensations, - ) = fully_fused_projection( - means, - None, - quats, - scales, - viewmats, - Ks, - width, - height, - packed=True, - sparse_grad=sparse_grad, - calc_compensations=calc_compensations, - camera_model=camera_model, + res = gsplat.fully_fused_projection( + means, None, quats, scales, viewmats, Ks, width, height, packed=True, + sparse_grad=sparse_grad, calc_compensations=calc_compensations, camera_model=camera_model ) - _radii, _means2d, _depths, _conics, _compensations = fully_fused_projection( - means, - None, - quats, - scales, - viewmats, - Ks, - width, - height, - packed=False, - calc_compensations=calc_compensations, - camera_model=camera_model, + _radii, _means2d, _depths, _conics, _compensations = gsplat.fully_fused_projection( + means, None, quats, scales, viewmats, Ks, width, height, packed=False, + calc_compensations=calc_compensations, camera_model=camera_model ) else: - covars, _ = quat_scale_to_covar_preci(quats, scales, triu=True) # [..., N, 6] - ( - batch_ids, - camera_ids, - gaussian_ids, - radii, - means2d, - depths, - conics, - compensations, - ) = fully_fused_projection( - means, - covars, - None, - None, - viewmats, - Ks, - width, - height, - packed=True, - sparse_grad=sparse_grad, - calc_compensations=calc_compensations, - camera_model=camera_model, + covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=True) + res = gsplat.fully_fused_projection( + means, covars, None, None, viewmats, Ks, width, height, packed=True, + sparse_grad=sparse_grad, calc_compensations=calc_compensations, camera_model=camera_model ) - _radii, _means2d, _depths, _conics, _compensations = fully_fused_projection( - means, - covars, - None, - None, - viewmats, - Ks, - width, - height, - packed=False, - calc_compensations=calc_compensations, - camera_model=camera_model, + _radii, _means2d, _depths, _conics, _compensations = gsplat.fully_fused_projection( + means, covars, None, None, viewmats, Ks, width, height, packed=False, + calc_compensations=calc_compensations, 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)) + + batch_ids, camera_ids, gaussian_ids, radii, means2d, depths, conics, compensations = res + 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) torch.testing.assert_close(__depths[sel], _depths[sel], rtol=1e-4, atol=1e-4) torch.testing.assert_close(__conics[sel], _conics[sel], rtol=1e-4, atol=1e-4) if calc_compensations: - torch.testing.assert_close( - __compensations[sel], _compensations[sel], rtol=1e-4, atol=1e-3 - ) + torch.testing.assert_close(__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() - + (_conics * v_conics).sum(), - (viewmats, quats, scales, means), - retain_graph=True, + (_means2d * v_means2d).sum() + (_depths * v_depths).sum() + (_conics * v_conics).sum(), + (viewmats, quats, scales, means), 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(), - (viewmats, quats, scales, means), - retain_graph=True, + (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) @@ -444,41 +262,31 @@ 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) - - B = math.prod(batch_dims) - C, N = 3, 1000 - I = B * C - width, height = 40, 60 + from gsplat._torch_impl import _isect_offset_encode, _isect_tiles + 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, "radii": torch.randint(0, width, (C, N, 2), device=device, dtype=torch.int32), "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( - means2d, radii, depths, tile_size, tile_width, tile_height - ) - isect_offsets = isect_offset_encode(isect_ids, I, tile_width, tile_height) + tiles_per_gauss, isect_ids, flatten_ids = gsplat.isect_tiles(means2d, radii, depths, tile_size, 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 - ) + _tiles_per_gauss, _isect_ids, _gauss_ids = _isect_tiles(means2d, radii, depths, tile_size, tile_width, tile_height) _isect_offsets = _isect_offset_encode(_isect_ids, I, tile_width, tile_height) torch.testing.assert_close(tiles_per_gauss, _tiles_per_gauss) @@ -487,59 +295,33 @@ 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( - { - "colors": torch.rand(C, N, channels, device=device), - "backgrounds": torch.rand((C, channels), device=device), - } - ) + test_data.update({ + "colors": torch.rand(C, N, channels, device=device), + "backgrounds": torch.rand((C, channels), device=device), + }) 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"] + 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"] - covars, _ = quat_scale_to_covar_preci(quats, scales, compute_preci=False, triu=True) - - # Project Gaussians to 2D - radii, means2d, depths, conics, compensations = fully_fused_projection( - means, covars, None, None, viewmats, Ks, width, height - ) + 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( - 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)) + 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 = 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 @@ -547,93 +329,28 @@ 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( - means2d, - conics, - colors, - opacities, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - backgrounds=backgrounds, + render_colors, render_alphas = gsplat.rasterize_to_pixels( + means2d, conics, colors, opacities, width, height, tile_size, + isect_offsets, 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, + 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) - - 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), - ) - ( - _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_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_coeffs, _v_dirs = torch.autograd.grad( - (_colors * v_colors).sum(), (coeffs, dirs), retain_graph=True, allow_unused=True + _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(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) + 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) From dedfc2da1ff5e50d1ff53258733add2c2255a76f Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Fri, 12 Sep 2025 18:08:24 -0700 Subject: [PATCH 02/56] Resolve Windows compiler / linker errors. fix mismatch between pass by value (CUDA) and pass by ref (sycl) for at::Tensor --- gsplat/sycl/CMakeLists.txt | 69 +-- gsplat/sycl/src/intersect_offset.cpp | 2 +- gsplat/sycl/src/intersect_tile.cpp | 10 +- gsplat/sycl/src/projection_ewa_simple_bwd.cpp | 10 +- gsplat/sycl/src/projection_ewa_simple_fwd.cpp | 6 +- .../src/quat_scale_to_covar_preci_bwd.cpp | 8 +- .../src/quat_scale_to_covar_preci_fwd.cpp | 4 +- .../sycl/src/rasterize_to_pixels_3dgs_bwd.cpp | 24 +- .../sycl/src/rasterize_to_pixels_3dgs_fwd.cpp | 16 +- gsplat/sycl/src/spherical_harmonics_bwd.cpp | 8 +- gsplat/sycl/src/spherical_harmonics_fwd.cpp | 6 +- setup.py | 45 +- tests/test_basic.py | 404 ++++++++++++++---- 13 files changed, 430 insertions(+), 182 deletions(-) diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt index 0e2373e0..aa8d0e4b 100644 --- a/gsplat/sycl/CMakeLists.txt +++ b/gsplat/sycl/CMakeLists.txt @@ -38,52 +38,53 @@ else() message(FATAL_ERROR "Could not find Torch via Python introspection. " "Please ensure PyTorch is installed or set CMAKE_PREFIX_PATH/Torch_DIR manually.") endif() +string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_SHARED_LIBRARY_SUFFIX}) execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], 'libtorch_python.so'))" + COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], '${TORCH_PYTHON_LIB_NAME}'))" OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE TORCH_PYTHON_LIB ) if (NOT EXISTS "${TORCH_PYTHON_LIB}") - message(FATAL_ERROR "Could not find libtorch_python.so at ${TORCH_PYTHON_LIB}. Please check your PyTorch installation.") + message(FATAL_ERROR "Could not find ${TORCH_PYTHON_LIB_NAME} at ${TORCH_PYTHON_LIB}. Please check your PyTorch installation.") else() message(STATUS "Found torch_python library at: ${TORCH_PYTHON_LIB}") endif() set(PYBIND11_FINDPYTHON ON) -find_package(pybind11 CONFIG REQUIRED) - -set( SYCL_SOURCES - ${CMAKE_CURRENT_SOURCE_DIR}/ext.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/adam.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/intersect_offset.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/intersect_tile.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/null.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_fused_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_fused_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_packed_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_2dgs_packed_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_fused_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_fused_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_packed_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_3dgs_packed_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_simple_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ewa_simple_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/projection_ut_3dgs_fused.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/quat_scale_to_covar_preci_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/quat_scale_to_covar_preci_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_indices_2dgs.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_indices_3dgs.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_2dgs_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_2dgs_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_3dgs_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_3dgs_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/relocation.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/spherical_harmonics_bwd.cpp - ${CMAKE_CURRENT_SOURCE_DIR}/src/spherical_harmonics_fwd.cpp +find_package(pybind11 CONFIG REQUIRED HINTS "${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) diff --git a/gsplat/sycl/src/intersect_offset.cpp b/gsplat/sycl/src/intersect_offset.cpp index 622001f0..3a50e3b0 100644 --- a/gsplat/sycl/src/intersect_offset.cpp +++ b/gsplat/sycl/src/intersect_offset.cpp @@ -9,7 +9,7 @@ namespace gsplat::xpu { at::Tensor intersect_offset( - const at::Tensor& isect_ids, // [n_isects] + const at::Tensor isect_ids, // [n_isects] const uint32_t I, const uint32_t tile_width, const uint32_t tile_height diff --git a/gsplat/sycl/src/intersect_tile.cpp b/gsplat/sycl/src/intersect_tile.cpp index 1fda9580..d40a6c9d 100644 --- a/gsplat/sycl/src/intersect_tile.cpp +++ b/gsplat/sycl/src/intersect_tile.cpp @@ -8,11 +8,11 @@ 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 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, diff --git a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp index a3a5e4b2..431d4521 100644 --- a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp @@ -7,14 +7,14 @@ 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 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] + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_covars2d // [..., C, N, 2, 2] ) { CHECK_CONTIGUOUS(means); diff --git a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp index 7179babe..1deafcbe 100644 --- a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp @@ -7,9 +7,9 @@ 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 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 diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp index a023b498..2517d04d 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -6,11 +6,11 @@ namespace gsplat::xpu { std::tuple quat_scale_to_covar_preci_bwd( - const at::Tensor& quats, // [..., 4] - const at::Tensor& scales, // [..., 3] + 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] + const at::optional v_covars, // [..., 3, 3] or [..., 6] + const at::optional v_precis // [..., 3, 3] or [..., 6] ) { CHECK_CONTIGUOUS(quats); CHECK_CONTIGUOUS(scales); diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp index 9eb92045..c457f2a9 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -6,8 +6,8 @@ namespace gsplat::xpu { std::tuple quat_scale_to_covar_preci_fwd( - const at::Tensor& quats, // [..., 4] - const at::Tensor& scales, // [..., 3] + const at::Tensor quats, // [..., 4] + const at::Tensor scales, // [..., 3] const bool compute_covar, const bool compute_preci, const bool triu diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp index de030b6d..1fd14693 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp @@ -103,25 +103,25 @@ void launch_rasterize_bwd_kernel( std::tuple rasterize_to_pixels_3dgs_bwd( // 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, + 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, + const at::Tensor tile_offsets, + const at::Tensor flatten_ids, // forward outputs - const at::Tensor& render_alphas, - const at::Tensor& last_ids, + 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, + const at::Tensor v_render_colors, + const at::Tensor v_render_alphas, // options bool absgrad ) { diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp index 953b0e72..a996ddfc 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp @@ -77,19 +77,19 @@ void launch_rasterize_kernel( std::tuple rasterize_to_pixels_3dgs_fwd( // 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, + 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 + const at::Tensor tile_offsets, + const at::Tensor flatten_ids ) { CHECK_CONTIGUOUS(means2d); CHECK_CONTIGUOUS(conics); diff --git a/gsplat/sycl/src/spherical_harmonics_bwd.cpp b/gsplat/sycl/src/spherical_harmonics_bwd.cpp index bc220c16..7845c74a 100644 --- a/gsplat/sycl/src/spherical_harmonics_bwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_bwd.cpp @@ -8,10 +8,10 @@ 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] + 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 ) { CHECK_CONTIGUOUS(dirs); diff --git a/gsplat/sycl/src/spherical_harmonics_fwd.cpp b/gsplat/sycl/src/spherical_harmonics_fwd.cpp index 83c66f9d..86385de2 100644 --- a/gsplat/sycl/src/spherical_harmonics_fwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_fwd.cpp @@ -7,9 +7,9 @@ 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 // [...] + const at::Tensor dirs, // [..., 3] + const at::Tensor coeffs, // [..., K, 3] + const at::optional masks // [...] ) { TORCH_CHECK(dirs.is_contiguous(), "Input 'dirs' tensor must be contiguous."); TORCH_CHECK(coeffs.is_contiguous(), "Input 'coeffs' tensor must be contiguous."); diff --git a/setup.py b/setup.py index c7fd6737..c258ea91 100644 --- a/setup.py +++ b/setup.py @@ -4,9 +4,8 @@ import pathlib import platform import sys - from setuptools import find_packages, setup -import subprocess +import subprocess as sp __version__ = None exec(open("gsplat/version.py", "r").read()) @@ -16,22 +15,26 @@ has_cuda = False try: import torch + has_cuda = torch.cuda.is_available() except ImportError: pass has_xpu = False -if not has_cuda: +if not has_cuda: try: import torch + has_xpu = torch.xpu.is_available() except (ImportError, AttributeError): - pass + pass has_sycl_compiler = False -if os.system('icpx --version > /dev/null 2>&1') == 0: - has_sycl_compiler = True -elif os.system('dpcpp --version > /dev/null 2>&1') == 0: +if ( + sp.run(["icpx", "--version"], stdout=sp.DEVNULL, stderr=sp.DEVNULL).returncode == 0 +) or ( + sp.run(["dpcpp", "--version"], stdout=sp.DEVNULL, stderr=sp.DEVNULL).returncode == 0 +): has_sycl_compiler = True BUILD_SYCL = has_xpu and has_sycl_compiler @@ -49,10 +52,12 @@ from torch.utils.cpp_extension import BuildExtension + class SyclBuildExtension(BuildExtension): """ Custom build class to orchestrate a CMake build for the SYCL backend. """ + def run(self): print("--- Running SYCL build via CMake ---") sycl_dir = os.path.abspath("gsplat/sycl") @@ -62,15 +67,23 @@ def run(self): install_dir = os.path.abspath(self.build_lib) - subprocess.check_call( - ["cmake", f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={os.path.join(install_dir, 'gsplat')}", sycl_dir], - cwd=build_dir + sp.check_call( + [ + "cmake", + "-G", + "Ninja", + "-DCMAKE_BUILD_TYPE=Release", + f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={os.path.join(install_dir, 'gsplat')}", + sycl_dir, + ], + cwd=build_dir, ) - subprocess.check_call( - ["cmake", "--build", ".", "--config", "Release", "--", f"-j{jobs}"], - cwd=build_dir + sp.check_call( + ["cmake", "--build", ".", "--config", "Release", "--", "-v", f"-j{jobs}"], + cwd=build_dir, ) + def get_ext(): from torch.utils.cpp_extension import BuildExtension @@ -157,10 +170,11 @@ def get_extensions(): cmdclass = {} packages_to_find = find_packages() from setuptools import Extension + if BUILD_SYCL: print("--- Configuring for SYCL build ---") cmdclass = {"build_ext": SyclBuildExtension} - ext_modules.append(Extension("gsplat.gsplat_sycl_kernels", sources=[])) + ext_modules.append(Extension("gsplat.gsplat_sycl_kernels", sources=[])) elif not BUILD_NO_CUDA: print("--- Configuring for CUDA build ---") ext_modules = get_extensions() @@ -175,7 +189,7 @@ def get_extensions(): keywords="gaussian, splatting, cuda, sycl", url=URL, download_url=f"{URL}/archive/gsplat-{__version__}.tar.gz", - python_requires=">=3.8", # Updated to match your CMake + python_requires=">=3.8", # Updated to match your CMake install_requires=[ "ninja", "numpy", @@ -202,7 +216,6 @@ def get_extensions(): cmdclass=cmdclass, packages=packages_to_find, include_package_data=True, - zip_safe=False, ) diff --git a/tests/test_basic.py b/tests/test_basic.py index a70db2b5..52a01528 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -23,10 +23,14 @@ elif gsplat.BACKEND == "cuda": device = torch.device("cuda:0") else: - device = torch.device("cpu") + device = None -requires_backend = pytest.mark.skipif(gsplat.BACKEND == "", reason="No CUDA or SYCL backend available") -requires_cuda = pytest.mark.skipif(gsplat.BACKEND != "cuda", reason="Test requires CUDA backend") +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="Test requires CUDA backend" +) def expand(data: dict, batch_dims: Tuple[int, ...]): @@ -76,7 +80,7 @@ def test_data(): def test_quat_scale_to_covar_preci(test_data, triu: bool, batch_dims: Tuple[int, ...]): 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"] @@ -90,8 +94,12 @@ def test_quat_scale_to_covar_preci(test_data, triu: bool, batch_dims: Tuple[int, 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)) - _v_quats, _v_scales = torch.autograd.grad((_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) + ) + _v_quats, _v_scales = torch.autograd.grad( + (_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) @@ -100,30 +108,50 @@ def test_quat_scale_to_covar_preci(test_data, triu: bool, batch_dims: Tuple[int, @pytest.mark.parametrize("camera_model", ["pinhole", "ortho", "fisheye"]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) 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._torch_impl import ( + _fisheye_proj, + _ortho_proj, + _persp_proj, + _world_to_cam, + ) torch.manual_seed(42) test_data = expand(test_data, batch_dims) - Ks, viewmats, height, width = test_data["Ks"], test_data["viewmats"], test_data["height"], test_data["width"] + Ks, viewmats, height, width = ( + test_data["Ks"], + test_data["viewmats"], + test_data["height"], + test_data["width"], + ) - covars, _ = gsplat.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 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": _means2d, _covars2d = _fisheye_proj(means, covars, Ks, width, height) - elif camera_model == "pinhole": _means2d, _covars2d = _persp_proj(means, covars, Ks, width, height) - else: assert_never(camera_model) + if camera_model == "ortho": + _means2d, _covars2d = _ortho_proj(means, covars, Ks, width, height) + elif camera_model == "fisheye": + _means2d, _covars2d = _fisheye_proj(means, covars, Ks, width, height) + elif camera_model == "pinhole": + _means2d, _covars2d = _persp_proj(means, covars, Ks, width, height) + else: + assert_never(camera_model) torch.testing.assert_close(means2d, _means2d, rtol=1e-4, atol=1e-4) torch.testing.assert_close(covars2d, _covars2d, rtol=1e-1, atol=3e-2) 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)) - _v_means, _v_covars = torch.autograd.grad((_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) + ) + _v_means, _v_covars = torch.autograd.grad( + (_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) @@ -133,13 +161,24 @@ def test_proj(test_data, camera_model: str, batch_dims: Tuple[int, ...]): @pytest.mark.parametrize("fused", [False, True]) @pytest.mark.parametrize("calc_compensations", [True, False]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) -def test_projection(test_data, fused: bool, calc_compensations: bool, camera_model: str, batch_dims: Tuple[int, ...]): +def test_projection( + test_data, + fused: bool, + calc_compensations: bool, + camera_model: str, + batch_dims: Tuple[int, ...], +): from gsplat._torch_impl import _fully_fused_projection torch.manual_seed(42) test_data = expand(test_data, batch_dims) - Ks, viewmats, height, width = test_data["Ks"], test_data["viewmats"], test_data["height"], test_data["width"] + 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 @@ -148,20 +187,42 @@ def test_projection(test_data, fused: bool, calc_compensations: bool, camera_mod if fused: radii, means2d, depths, conics, compensations = gsplat.fully_fused_projection( - means, None, quats, scales, viewmats, Ks, width, height, - calc_compensations=calc_compensations, camera_model=camera_model + means, + None, + quats, + scales, + viewmats, + Ks, + width, + height, + calc_compensations=calc_compensations, + camera_model=camera_model, ) else: covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=True) radii, means2d, depths, conics, compensations = gsplat.fully_fused_projection( - means, covars, None, None, viewmats, Ks, width, height, - calc_compensations=calc_compensations, camera_model=camera_model + means, + covars, + None, + None, + viewmats, + Ks, + width, + height, + calc_compensations=calc_compensations, + camera_model=camera_model, ) - + _covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=False) _radii, _means2d, _depths, _conics, _compensations = _fully_fused_projection( - means, _covars, viewmats, Ks, width, height, - calc_compensations=calc_compensations, camera_model=camera_model + means, + _covars, + viewmats, + Ks, + width, + height, + calc_compensations=calc_compensations, + camera_model=camera_model, ) valid = (radii > 0).all(dim=-1) & (_radii > 0).all(dim=-1) @@ -170,15 +231,37 @@ def test_projection(test_data, fused: bool, calc_compensations: bool, camera_mod torch.testing.assert_close(depths[valid], _depths[valid], rtol=1e-4, atol=1e-4) torch.testing.assert_close(conics[valid], _conics[valid], rtol=1e-4, atol=1e-4) if calc_compensations: - torch.testing.assert_close(compensations[valid], _compensations[valid], rtol=1e-4, atol=1e-3) + torch.testing.assert_close( + compensations[valid], _compensations[valid], rtol=1e-4, atol=1e-3 + ) - 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) - 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) - _v_viewmats, _v_quats, _v_scales, _v_means = torch.autograd.grad(_grad_sum, (viewmats, quats, scales, means)) + 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) + ) + 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) + ) + _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) torch.testing.assert_close(v_quats, _v_quats, rtol=2e-1, atol=2e-2) @@ -192,11 +275,23 @@ def test_projection(test_data, fused: bool, calc_compensations: bool, camera_mod @pytest.mark.parametrize("calc_compensations", [False, True]) @pytest.mark.parametrize("camera_model", ["pinhole", "ortho", "fisheye"]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) -def test_fully_fused_projection_packed(test_data, fused: bool, sparse_grad: bool, calc_compensations: bool, camera_model: str, batch_dims: Tuple[int, ...]): - +def test_fully_fused_projection_packed( + test_data, + fused: bool, + sparse_grad: bool, + calc_compensations: bool, + camera_model: str, + batch_dims: Tuple[int, ...], +): + torch.manual_seed(42) test_data = expand(test_data, batch_dims) - Ks, viewmats, height, width = test_data["Ks"], test_data["viewmats"], test_data["height"], test_data["width"] + 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 @@ -205,36 +300,107 @@ def test_fully_fused_projection_packed(test_data, fused: bool, sparse_grad: bool if fused: res = gsplat.fully_fused_projection( - means, None, quats, scales, viewmats, Ks, width, height, packed=True, - sparse_grad=sparse_grad, calc_compensations=calc_compensations, camera_model=camera_model + means, + None, + quats, + scales, + viewmats, + Ks, + width, + height, + packed=True, + sparse_grad=sparse_grad, + calc_compensations=calc_compensations, + camera_model=camera_model, ) - _radii, _means2d, _depths, _conics, _compensations = gsplat.fully_fused_projection( - means, None, quats, scales, viewmats, Ks, width, height, packed=False, - calc_compensations=calc_compensations, camera_model=camera_model + _radii, _means2d, _depths, _conics, _compensations = ( + gsplat.fully_fused_projection( + means, + None, + quats, + scales, + viewmats, + Ks, + width, + height, + packed=False, + calc_compensations=calc_compensations, + camera_model=camera_model, + ) ) else: covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=True) res = gsplat.fully_fused_projection( - means, covars, None, None, viewmats, Ks, width, height, packed=True, - sparse_grad=sparse_grad, calc_compensations=calc_compensations, camera_model=camera_model + means, + covars, + None, + None, + viewmats, + Ks, + width, + height, + packed=True, + sparse_grad=sparse_grad, + calc_compensations=calc_compensations, + camera_model=camera_model, ) - _radii, _means2d, _depths, _conics, _compensations = gsplat.fully_fused_projection( - means, covars, None, None, viewmats, Ks, width, height, packed=False, - calc_compensations=calc_compensations, camera_model=camera_model + _radii, _means2d, _depths, _conics, _compensations = ( + gsplat.fully_fused_projection( + means, + covars, + None, + None, + viewmats, + Ks, + width, + height, + packed=False, + calc_compensations=calc_compensations, + camera_model=camera_model, + ) ) - - batch_ids, camera_ids, gaussian_ids, radii, means2d, depths, conics, compensations = res + + ( + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + compensations, + ) = res 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)) + __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(indices, compensations, sparse_shape).to_dense().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) @@ -242,19 +408,35 @@ def test_fully_fused_projection_packed(test_data, fused: bool, sparse_grad: bool torch.testing.assert_close(__depths[sel], _depths[sel], rtol=1e-4, atol=1e-4) torch.testing.assert_close(__conics[sel], _conics[sel], rtol=1e-4, atol=1e-4) if calc_compensations: - torch.testing.assert_close(__compensations[sel], _compensations[sel], rtol=1e-4, atol=1e-3) + torch.testing.assert_close( + __compensations[sel], _compensations[sel], rtol=1e-4, atol=1e-3 + ) - v_means2d, v_depths, v_conics = torch.randn_like(_means2d) * sel[..., None], torch.randn_like(_depths) * sel, 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() + (_conics * v_conics).sum(), - (viewmats, quats, scales, means), retain_graph=True + (_means2d * v_means2d).sum() + + (_depths * v_depths).sum() + + (_conics * v_conics).sum(), + (viewmats, quats, scales, means), + retain_graph=True, ) v_viewmats, v_quats, v_scales, v_means = torch.autograd.grad( - (means2d * v_means2d[sel]).sum() + (depths * v_depths[sel]).sum() + (conics * v_conics[sel]).sum(), - (viewmats, quats, scales, means), retain_graph=True + (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_scales, v_means = v_quats.to_dense(), v_scales.to_dense(), 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) @@ -271,22 +453,32 @@ def test_isect(test_data, batch_dims: Tuple[int, ...]): 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, "radii": torch.randint(0, width, (C, N, 2), device=device, dtype=torch.int32), "depths": torch.rand(C, N, device=device), } test_data = expand(test_data, batch_dims) - means2d, radii, depths = test_data["means2d"], test_data["radii"], test_data["depths"] + means2d, radii, depths = ( + test_data["means2d"], + test_data["radii"], + test_data["depths"], + ) tile_size = 16 - tile_width, tile_height = math.ceil(width / tile_size), 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 = gsplat.isect_tiles(means2d, radii, depths, tile_size, tile_width, tile_height) + tiles_per_gauss, isect_ids, flatten_ids = gsplat.isect_tiles( + means2d, radii, depths, tile_size, 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) + _tiles_per_gauss, _isect_ids, _gauss_ids = _isect_tiles( + means2d, radii, depths, tile_size, tile_width, tile_height + ) _isect_offsets = _isect_offset_encode(_isect_ids, I, tile_width, tile_height) torch.testing.assert_close(tiles_per_gauss, _tiles_per_gauss) @@ -305,23 +497,45 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. 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({ - "colors": torch.rand(C, N, channels, device=device), - "backgrounds": torch.rand((C, channels), device=device), - }) + test_data.update( + { + "colors": torch.rand(C, N, channels, device=device), + "backgrounds": torch.rand((C, channels), device=device), + } + ) test_data = expand(test_data, batch_dims) - 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"] + 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"] - 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) + 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)) tile_size = 16 if channels <= 32 else 4 - 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 = gsplat.isect_offset_encode(isect_ids, I, tile_width, tile_height).reshape(batch_dims + (C, tile_height, tile_width)) + 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 = 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 @@ -330,24 +544,44 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. backgrounds.requires_grad = True render_colors, render_alphas = gsplat.rasterize_to_pixels( - means2d, conics, colors, opacities, width, height, tile_size, - isect_offsets, flatten_ids, backgrounds=backgrounds + means2d, + conics, + colors, + opacities, + width, + height, + tile_size, + isect_offsets, + 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 + 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_render_colors, v_render_alphas = torch.randn_like(render_colors), torch.randn_like(render_alphas) + 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) + (render_colors * v_render_colors).sum() + + (render_alphas * v_render_alphas).sum(), + (means2d, conics, colors, opacities, backgrounds), ) _grads = torch.autograd.grad( - (_render_colors * v_render_colors).sum() + (_render_alphas * v_render_alphas).sum(), - (means2d, conics, colors, opacities, backgrounds) + (_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) From 6f44999298276c54289ba8d2685629388a616a64 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Fri, 12 Sep 2025 22:56:38 -0700 Subject: [PATCH 03/56] Fix Windows linker errors --- gsplat/sycl/CMakeLists.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt index aa8d0e4b..194bdb2e 100644 --- a/gsplat/sycl/CMakeLists.txt +++ b/gsplat/sycl/CMakeLists.txt @@ -38,7 +38,7 @@ else() message(FATAL_ERROR "Could not find Torch via Python introspection. " "Please ensure PyTorch is installed or set CMAKE_PREFIX_PATH/Torch_DIR manually.") endif() -string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_SHARED_LIBRARY_SUFFIX}) +string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_LINK_LIBRARY_SUFFIX}) execute_process( COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], '${TORCH_PYTHON_LIB_NAME}'))" OUTPUT_STRIP_TRAILING_WHITESPACE @@ -104,7 +104,7 @@ 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) +target_link_libraries(${SYCL_MODULE_NAME} PRIVATE torch ${TORCH_PYTHON_LIB}) # Fix for icx: error: '-MP' is not supported with offloading enabled From a9eaf9977763d73eb4c27b8a737a4f9e260195e5 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Mon, 15 Sep 2025 00:30:09 -0700 Subject: [PATCH 04/56] Loading kernels now works in Windows. --- gsplat/sycl/_backend.py | 21 +++++++++++++++++++-- 1 file changed, 19 insertions(+), 2 deletions(-) diff --git a/gsplat/sycl/_backend.py b/gsplat/sycl/_backend.py index 216999ca..fbc8e5f0 100644 --- a/gsplat/sycl/_backend.py +++ b/gsplat/sycl/_backend.py @@ -1,7 +1,24 @@ _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") \ No newline at end of file +except ImportError: + raise ImportError("Unable to find compiled sycl kernels package") + +if os.name == "nt": + for dp in dllpath: + dp.close() From 2f28ba3a65e5fdb1052d891270abe859c8a36f17 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Tue, 23 Sep 2025 08:08:53 +0000 Subject: [PATCH 05/56] COrrected libtorch path issues --- gsplat/sycl/CMakeLists.txt | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt index 194bdb2e..28007357 100644 --- a/gsplat/sycl/CMakeLists.txt +++ b/gsplat/sycl/CMakeLists.txt @@ -38,13 +38,28 @@ else() message(FATAL_ERROR "Could not find Torch via Python introspection. " "Please ensure PyTorch is installed or set CMAKE_PREFIX_PATH/Torch_DIR manually.") endif() -string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_LINK_LIBRARY_SUFFIX}) + +string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_SHARED_LIBRARY_SUFFIX}) execute_process( COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], '${TORCH_PYTHON_LIB_NAME}'))" OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE TORCH_PYTHON_LIB ) +# Add a fallback to find the file if it exists with a different extension +if (NOT EXISTS "${TORCH_PYTHON_LIB}") + execute_process( + COMMAND "${Python_EXECUTABLE}" -c "import os, glob; from torch.utils import cpp_extension; base_path = os.path.join(cpp_extension.library_paths(True)[0], '${CMAKE_SHARED_LIBRARY_PREFIX}torch_python'); matches = glob.glob(base_path + '.*'); print(matches[0] if matches else '')" + OUTPUT_STRIP_TRAILING_WHITESPACE + OUTPUT_VARIABLE TORCH_PYTHON_LIB_FALLBACK + ) + if (EXISTS "${TORCH_PYTHON_LIB_FALLBACK}") + set(TORCH_PYTHON_LIB "${TORCH_PYTHON_LIB_FALLBACK}") + message(STATUS "Found torch_python library using fallback: ${TORCH_PYTHON_LIB}") + endif() +endif() + + if (NOT EXISTS "${TORCH_PYTHON_LIB}") message(FATAL_ERROR "Could not find ${TORCH_PYTHON_LIB_NAME} at ${TORCH_PYTHON_LIB}. Please check your PyTorch installation.") else() From 3fff3defebf7bc344c2303698f34487a0c2b87a8 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Wed, 24 Sep 2025 14:37:39 +0000 Subject: [PATCH 06/56] Updated correct shape calculation --- .../sycl/src/quat_scale_to_covar_preci_bwd.cpp | 2 +- .../sycl/src/quat_scale_to_covar_preci_fwd.cpp | 18 ++++++++++++++---- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp index 2517d04d..1868cd15 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -22,7 +22,7 @@ std::tuple quat_scale_to_covar_preci_bwd( } 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.size(0); + const int64_t N = quats.numel() / 4; at::Tensor v_quats = at::empty_like(quats); at::Tensor v_scales = at::empty_like(scales); diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp index c457f2a9..8ab9ce49 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -16,19 +16,29 @@ std::tuple quat_scale_to_covar_preci_fwd( CHECK_CONTIGUOUS(scales); TORCH_CHECK(compute_covar || compute_preci, "Must compute at least one of covar or preci"); - const int64_t N = quats.size(0); + 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 = triu ? at::empty({N, 6}, options) : at::empty({N, 3, 3}, options); + covars = at::empty(out_shape, options); } else { covars = at::empty({0}, options); } - at::Tensor precis; if (compute_preci) { - precis = triu ? at::empty({N, 6}, options) : at::empty({N, 3, 3}, options); + precis = at::empty(out_shape, options); } else { precis = at::empty({0}, options); } From 10c0efd7ccd78e6b0bc12007a8f37de948ce4874 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Wed, 24 Sep 2025 14:57:56 +0000 Subject: [PATCH 07/56] Updated proj changes --- gsplat/sycl/include/kernels/ProjBwdKernel.hpp | 12 +++++++----- gsplat/sycl/include/kernels/ProjFwdKernel.hpp | 10 +++++++--- gsplat/sycl/src/projection_ewa_simple_bwd.cpp | 10 +++++----- gsplat/sycl/src/projection_ewa_simple_fwd.cpp | 19 ++++++++++++++----- 4 files changed, 33 insertions(+), 18 deletions(-) diff --git a/gsplat/sycl/include/kernels/ProjBwdKernel.hpp b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp index ffe84fe0..9d363fa6 100644 --- a/gsplat/sycl/include/kernels/ProjBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp @@ -46,18 +46,20 @@ struct ProjBwdKernel{ void operator()(sycl::nd_item<1> work_item) const { uint32_t idx = work_item.get_global_id(0); - if (idx >= m_C * m_N) { + 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 uint32_t cid = idx / m_N; // camera id - - // shift pointers to the current camera and gaussian 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); - const T* Ks = m_Ks + (cid * 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); diff --git a/gsplat/sycl/include/kernels/ProjFwdKernel.hpp b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp index c50b6229..8cef2b28 100644 --- a/gsplat/sycl/include/kernels/ProjFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp @@ -40,14 +40,18 @@ struct ProjFwdKernel{ void operator()(sycl::nd_item<1> work_item) const { uint32_t idx = work_item.get_global_id(0); - if (idx >= m_C * m_N) { + 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 cid = idx / m_N; // camera id + 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 + (cid * 9); + const T* Ks = m_Ks + (bid * m_C * 9) + (cid * 9); + T* means2d = m_means2d + (idx * 2); T* covars2d = m_covars2d + (idx * 4); diff --git a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp index 431d4521..8b1deead 100644 --- a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp @@ -26,17 +26,17 @@ std::tuple projection_ewa_simple_bwd( 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); - at::Tensor v_means = at::empty({C, N, 3}, means.options()); - at::Tensor v_covars = at::empty({C, N, 3, 3}, covars.options()); - - if (C > 0 && N > 0) { + if (total_gaussians > 0) { auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + size_t numWorkGrps = (total_gaussians + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; - size_t numWorkGrps = (C * 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); diff --git a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp index 1deafcbe..00441691 100644 --- a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp @@ -23,15 +23,24 @@ std::tuple projection_ewa_simple_fwd( const uint32_t C = means.size(-3); const uint32_t N = means.size(-2); + const uint32_t total_gaussians = means.numel() / 3; - at::Tensor means2d = at::empty({C, N, 2}, means.options()); - at::Tensor covars2d = at::empty({C, N, 2, 2}, covars.options()); + auto options = means.options(); + at::DimVector batch_dims(means.sizes().slice(0, means.dim() - 3)); - if (C > 0 && N > 0) { + 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; - - size_t numWorkGrps = (C * 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); From 7b3bfb8c8a7a809845c4c739f8be724567b4e2e0 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Wed, 24 Sep 2025 15:18:08 +0000 Subject: [PATCH 08/56] corrected isect code --- .../sycl/include/kernels/IsectTilesKernel.hpp | 22 +++++++++++-------- gsplat/sycl/src/intersect_tile.cpp | 10 ++++----- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/gsplat/sycl/include/kernels/IsectTilesKernel.hpp b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp index ba6bb81a..69f03f3d 100644 --- a/gsplat/sycl/include/kernels/IsectTilesKernel.hpp +++ b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp @@ -80,8 +80,9 @@ struct IsectTilesKernel { return; } - const T radius = m_radii[idx]; - if (radius <= 0) { + 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; } @@ -90,16 +91,17 @@ struct IsectTilesKernel { vec2 mean2d = glm::make_vec2(m_means2d + 2 * idx); - T tile_radius = radius / static_cast(m_tile_size); + 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; - tile_min.x = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::floor(tile_x - tile_radius)), m_tile_width); - tile_min.y = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::floor(tile_y - tile_radius)), m_tile_height); - - tile_max.x = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::ceil(tile_x + tile_radius)), m_tile_width); - tile_max.y = sycl::min( sycl::max((uint32_t)0, (uint32_t)sycl::ceil(tile_y + tile_radius)), m_tile_height); + // 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 @@ -122,7 +124,9 @@ struct IsectTilesKernel { const int64_t cid_enc = cid << (32 + m_tile_n_bits); - int64_t depth_id_enc = (int64_t) * (int32_t *)&(m_depths[idx]); + 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) { diff --git a/gsplat/sycl/src/intersect_tile.cpp b/gsplat/sycl/src/intersect_tile.cpp index d40a6c9d..0c0e051f 100644 --- a/gsplat/sycl/src/intersect_tile.cpp +++ b/gsplat/sycl/src/intersect_tile.cpp @@ -38,7 +38,7 @@ std::tuple intersect_tile( 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(1); + N = means2d.size(-2); total_elems = C * N; } @@ -49,8 +49,8 @@ std::tuple intersect_tile( at::empty({0}, at::kInt) ); } - - at::Tensor tiles_per_gauss = at::empty_like(depths, depths.options().dtype(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; @@ -86,8 +86,8 @@ std::tuple intersect_tile( n_isects = cum_tiles_per_gauss.slice(0, -1).item(); } - at::Tensor isect_ids = at::empty({n_isects}, at::kLong); - at::Tensor flatten_ids = at::empty({n_isects}, at::kInt); + 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) { From 0d25c9e04cee8ea54da2bdb9c10407ed5f857369 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Thu, 25 Sep 2025 06:35:59 +0000 Subject: [PATCH 09/56] Update fully fused projection kernels --- .../kernels/FullyFusedProjectionBwdKernel.hpp | 99 +++++++++------ .../kernels/FullyFusedProjectionFwdKernel.hpp | 120 ++++++++++++------ gsplat/sycl/include/utils.hpp | 2 +- .../src/projection_ewa_3dgs_fused_bwd.cpp | 76 ++++++++++- .../src/projection_ewa_3dgs_fused_fwd.cpp | 81 +++++++++++- 5 files changed, 288 insertions(+), 90 deletions(-) diff --git a/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp index b3f72703..70e210e3 100644 --- a/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp @@ -13,35 +13,40 @@ 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; // [N, 3] - const T* m_covars; // [N, 6] optional - const T* m_quats; // [N, 4] optional - const T* m_scales; // [N, 3] optional - const T* m_viewmats; // [C, 4, 4] - const T* m_Ks; // [C, 3, 3] + 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 - const int32_t* m_radii; // [C, N] - const T* m_conics; // [C, N, 3] - const T* m_compensations; // [C, N] optional + // 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; // [C, N, 2] - const T* m_v_depths; // [C, N] - const T* m_v_conics; // [C, N, 3] - const T* m_v_compensations; // [C, N] optional + 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; // [N, 3] - T* m_v_covars; // [N, 6] optional - T* m_v_quats; // [N, 4] optional - T* m_v_scales; // [N, 3] optional - T* m_v_viewmats;// [C, 4, 4] optional - + 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, @@ -67,31 +72,35 @@ struct FullyFusedProjectionBwdKernel{ T* v_scales, T* v_viewmats ) - : m_C(C), m_N(N), m_means(means), m_covars(covars), m_quats(quats), m_scales(scales), + // 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 + void operator()(sycl::nd_item<1> work_item) const { uint32_t idx = work_item.get_global_id(0); - if (idx >= m_C * m_N || m_radii[idx] <= 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; } - - const uint32_t cid = idx / m_N; // camera id + + // 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 - - // shift pointers to the current camera and gaussian - const T* means = m_means + (gid * 3); - const T* viewmats = m_viewmats + (cid * 16); - const T* Ks = m_Ks + (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); + + // 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]); @@ -127,7 +136,8 @@ struct FullyFusedProjectionBwdKernel{ vec4 quat; vec3 scale; if (m_covars != nullptr) { - const T* covars = m_covars + (gid * 6); + // Changed: Updated pointer arithmetic + const T* covars = m_covars + bid * m_N * 6 + gid * 6; covar = mat3( covars[0], covars[1], @@ -141,8 +151,9 @@ struct FullyFusedProjectionBwdKernel{ ); } else { // compute from quaternions and scales - quat = glm::make_vec4(m_quats + (gid * 4)); - scale = glm::make_vec3(m_scales + (gid * 3)); + // 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; @@ -220,7 +231,8 @@ struct FullyFusedProjectionBwdKernel{ 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); + // 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]); @@ -228,7 +240,8 @@ struct FullyFusedProjectionBwdKernel{ } if (m_v_covars != nullptr) { - T* v_covars = m_v_covars + (gid * 6); + // 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]); @@ -243,8 +256,9 @@ struct FullyFusedProjectionBwdKernel{ quat_scale_to_covar_vjp( quat, scale, rotmat, v_covar, v_quat, v_scale ); - T* v_quats = m_v_quats + (gid * 4); - T* v_scales = m_v_scales + (gid * 3); + // 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]); @@ -255,7 +269,8 @@ struct FullyFusedProjectionBwdKernel{ } if (m_v_viewmats != nullptr) { - T* v_viewmats = m_v_viewmats + (cid * 16); + // 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 @@ -270,4 +285,4 @@ struct FullyFusedProjectionBwdKernel{ #endif //FullyFusedProjectionBwdKernel_HPP -} // namespace gsplat::xpu +} // 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 index d248d14b..e5074e46 100644 --- a/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp @@ -11,14 +11,18 @@ 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; // [N, 3] - const T* m_covars; // [N, 6] optional - const T* m_quats; // [N, 4] optional - const T* m_scales; // [N, 3] optional - const T* m_viewmats; // [C, 4, 4] - const T* m_Ks; // [C, 3, 3] + 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; @@ -27,19 +31,24 @@ struct FullyFusedProjectionFwdKernel{ const T m_radius_clip; const CameraModelType m_camera_model; // outputs - int32_t * m_radii; // [C, N] - T* m_means2d; // [C, N, 2] - T* m_depths; // [C, N] - T* m_conics; // [C, N, 3] - T* m_compensations; // [C, N] optional + // 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, @@ -55,25 +64,30 @@ struct FullyFusedProjectionFwdKernel{ T* conics, T* compensations ) - : 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), + // 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) + 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 + void operator()(sycl::nd_item<1> work_item) const { uint32_t idx = work_item.get_global_id(0); - if (idx >= m_C * m_N) { + // Changed: Updated upper bound to include B + if (idx >= m_B * m_C * m_N) { return; } - const uint32_t cid = idx / m_N; // camera id + // 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 - const T* means = m_means + (gid * 3); - const T* viewmats = m_viewmats + (cid * 16); - const T* Ks = m_Ks + (cid * 9); + // 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( @@ -88,19 +102,22 @@ struct FullyFusedProjectionFwdKernel{ 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) { - m_radii[idx] = 0; + // 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) { - const T* covars = m_covars + (gid * 6); + // Changed: Updated pointer arithmetic + const T* covars = m_covars + bid * m_N * 6 + gid * 6; covar = mat3( covars[0], covars[1], @@ -114,15 +131,16 @@ struct FullyFusedProjectionFwdKernel{ ); } else { // compute from quaternions and scales - const T* quats = m_quats + (gid * 4); - const T* scales = m_scales + (gid * 3); + // 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; @@ -175,7 +193,9 @@ struct FullyFusedProjectionFwdKernel{ T compensation; T det = add_blur(m_eps2d, covar2d, compensation); if (det <= 0.f) { - m_radii[idx] = 0; + // Changed: Set both radii to 0 + m_radii[idx * 2] = 0; + m_radii[idx * 2 + 1] = 0; return; } @@ -183,25 +203,43 @@ struct FullyFusedProjectionFwdKernel{ mat2 covar2d_inv; inverse(covar2d, covar2d_inv); - // take 3 sigma as the radius (non differentiable) - T b = 0.5f * (covar2d[0][0] + covar2d[1][1]); - T v1 = b + sycl::sqrt(sycl::max(0.01f, b * b - det)); - T radius = sycl::ceil(3.f * sycl::sqrt(v1)); + // 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))); + } - if (radius <= m_radius_clip) { - m_radii[idx] = 0; + 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 <= 0 || mean2d.x - radius >= m_image_width || - mean2d.y + radius <= 0 || mean2d.y - radius >= m_image_height) { - m_radii[idx] = 0; + 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 - m_radii[idx] = (int32_t)radius; + // 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; @@ -211,9 +249,9 @@ struct FullyFusedProjectionFwdKernel{ if (m_compensations != nullptr) { m_compensations[idx] = compensation; } - + } - + }; #endif //FullyFusedProjectionFwdKernel_HPP diff --git a/gsplat/sycl/include/utils.hpp b/gsplat/sycl/include/utils.hpp index 6d63b3ac..094fb354 100644 --- a/gsplat/sycl/include/utils.hpp +++ b/gsplat/sycl/include/utils.hpp @@ -40,7 +40,7 @@ inline T add_blur(const T eps2d, mat2 &covar, T &compensation) { 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(0.f, det_orig / det_blur)); + compensation = sycl::sqrt(sycl::max(static_cast(0), det_orig / det_blur)); return det_blur; } diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp index be768bf6..6f0be199 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp @@ -1,11 +1,11 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/FullyFusedProjectionBwdKernel.hpp" namespace gsplat::xpu { - + std::tuple projection_ewa_3dgs_fused_bwd( // fwd inputs @@ -30,7 +30,77 @@ projection_ewa_3dgs_fused_bwd( const at::optional v_compensations, // [..., C, N] optional const bool viewmats_requires_grad ) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + // Input validation + CHECK_CONTIGUOUS(means); + if (covars.has_value()) CHECK_CONTIGUOUS(covars.value()); + if (quats.has_value()) CHECK_CONTIGUOUS(quats.value()); + if (scales.has_value()) CHECK_CONTIGUOUS(scales.value()); + CHECK_CONTIGUOUS(viewmats); + CHECK_CONTIGUOUS(Ks); + CHECK_CONTIGUOUS(radii); + CHECK_CONTIGUOUS(conics); + if (compensations.has_value()) CHECK_CONTIGUOUS(compensations.value()); + CHECK_CONTIGUOUS(v_means2d); + CHECK_CONTIGUOUS(v_depths); + CHECK_CONTIGUOUS(v_conics); + if (v_compensations.has_value()) CHECK_CONTIGUOUS(v_compensations.value()); + + // 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 index c1eb4d05..1e42afee 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp @@ -1,11 +1,11 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/FullyFusedProjectionFwdKernel.hpp" namespace gsplat::xpu { - + std::tuple< at::Tensor, at::Tensor, @@ -29,7 +29,82 @@ projection_ewa_3dgs_fused_fwd( const bool calc_compensations, const CameraModelType camera_model ) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + CHECK_CONTIGUOUS(means); + CHECK_CONTIGUOUS(viewmats); + CHECK_CONTIGUOUS(Ks); + if (covars.has_value()) CHECK_CONTIGUOUS(covars.value()); + if (quats.has_value()) CHECK_CONTIGUOUS(quats.value()); + if (scales.has_value()) CHECK_CONTIGUOUS(scales.value()); + if (opacities.has_value()) CHECK_CONTIGUOUS(opacities.value()); + + 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); + const uint32_t C = viewmats.size(-3); + const uint32_t B = means.numel() / (N * 3); + 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 From 02a596ba92083af6c6530ffa65c19f5008d4df68 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Thu, 25 Sep 2025 07:59:27 +0000 Subject: [PATCH 10/56] wip packed kernel --- gsplat/sycl/include/Common.h | 2 +- .../kernels/PackedProjectionBwdKernel.hpp | 224 ++++++++++++++++ .../kernels/PackedProjectionFwdKernel.hpp | 246 ++++++++++++++++++ .../src/projection_ewa_3dgs_packed_bwd.cpp | 89 ++++++- .../src/projection_ewa_3dgs_packed_fwd.cpp | 149 ++++++++++- 5 files changed, 696 insertions(+), 14 deletions(-) create mode 100644 gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp create mode 100644 gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp diff --git a/gsplat/sycl/include/Common.h b/gsplat/sycl/include/Common.h index df56d625..407ae2ca 100644 --- a/gsplat/sycl/include/Common.h +++ b/gsplat/sycl/include/Common.h @@ -28,5 +28,5 @@ enum CameraModelType { }; #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/kernels/PackedProjectionBwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp new file mode 100644 index 00000000..3bf876bb --- /dev/null +++ b/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp @@ -0,0 +1,224 @@ +#ifndef PackedProjectionBwdKernel_HPP +#define PackedProjectionBwdKernel_HPP + +#include +#include "utils.hpp" +#include "quat.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "proj.hpp" +#include "transform.hpp" + +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) { + sycl::atomic_ref ref(v_means_out[i]); + ref.fetch_add(v_mean[i]); + } + } + if (m_v_covars != nullptr) { + T* v_covars_out = m_v_covars + bid * m_N * 6 + gid * 6; + sycl::atomic_ref(v_covars_out[0]).fetch_add(v_covar[0][0]); + sycl::atomic_ref(v_covars_out[1]).fetch_add(v_covar[0][1] + v_covar[1][0]); + sycl::atomic_ref(v_covars_out[2]).fetch_add(v_covar[0][2] + v_covar[2][0]); + sycl::atomic_ref(v_covars_out[3]).fetch_add(v_covar[1][1]); + sycl::atomic_ref(v_covars_out[4]).fetch_add(v_covar[1][2] + v_covar[2][1]); + sycl::atomic_ref(v_covars_out[5]).fetch_add(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) sycl::atomic_ref(v_quats_out[i]).fetch_add(v_quat[i]); + for (int i = 0; i < 3; ++i) sycl::atomic_ref(v_scales_out[i]).fetch_add(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 + sycl::atomic_ref ref(v_viewmats_out[i * 4 + j]); + ref.fetch_add(v_R[j][i]); + } + sycl::atomic_ref ref(v_viewmats_out[i * 4 + 3]); + ref.fetch_add(v_t[i]); + } + } + } +}; + +} // namespace gsplat::xpu + +#endif // PackedProjectionBwdKernel_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp new file mode 100644 index 00000000..23419234 --- /dev/null +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -0,0 +1,246 @@ +#ifndef PackedProjectionFwdKernel_HPP +#define PackedProjectionFwdKernel_HPP + +#include +#include "utils.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "proj.hpp" +#include "transform.hpp" + +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 int32_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 int32_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_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::min(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 --- + int32_t thread_data = static_cast(valid); + + if (m_block_cnts != nullptr) { + // First pass: Count visible Gaussians in this block. + // Check if any thread in the group has a valid Gaussian. + 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. + bool any_valid = sycl::any_of_group(group, valid); + if (any_valid) { + // Perform an exclusive scan to find the local offset for this thread. + int32_t local_offset = sycl::exclusive_scan_over_group(group, thread_data, sycl::plus<>()); + + if (valid) { + int32_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 + +#endif // PackedProjectionFwdKernel_HPP \ 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 index c2b6606c..2650aaca 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp @@ -1,11 +1,11 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/PackedProjectionBwdKernel.hpp" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple projection_ewa_3dgs_packed_bwd( // fwd inputs @@ -31,9 +31,84 @@ projection_ewa_3dgs_packed_bwd( const at::Tensor v_conics, // [nnz, 3] const at::optional v_compensations, // [nnz] optional const bool viewmats_requires_grad, - const bool sparse_grad -) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + const bool sparse_grad) { + + TORCH_CHECK(means.is_contiguous(), "Input 'means' tensor must be contiguous."); + TORCH_CHECK(viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous."); + TORCH_CHECK(Ks.is_contiguous(), "Input 'Ks' tensor must be contiguous."); + TORCH_CHECK(batch_ids.is_contiguous(), "Input 'batch_ids' tensor must be contiguous."); + TORCH_CHECK(means.device().type() == at::kXPU, "Input tensors must be on XPU device."); + + 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 +} // 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 index 855e85e9..0c60c812 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -1,11 +1,11 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/PackedProjectionFwdKernel.hpp" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple< at::Tensor, at::Tensor, @@ -31,9 +31,146 @@ projection_ewa_3dgs_packed_fwd( const float far_plane, const float radius_clip, const bool calc_compensations, - const CameraModelType camera_model -) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + const CameraModelType camera_model) { + + TORCH_CHECK(means.is_contiguous(), "Input 'means' tensor must be contiguous."); + TORCH_CHECK(viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous."); + TORCH_CHECK(Ks.is_contiguous(), "Input 'Ks' tensor must be contiguous."); + TORCH_CHECK(means.device().type() == at::kXPU, "Input tensors must be on XPU device."); + + 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); + } + + // --- Start of Correction --- + // Changed block_cnts to kLong to satisfy at::cumsum requirements + at::Tensor block_cnts = at::empty({(long)n_blocks}, long_opts); + // --- End of Correction --- + + 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 + // --- Start of Correction --- + // The kernel expects int32_t*, but the tensor is int64_t. + // This cast is safe because the counts per block will not exceed int32_t max. + (int32_t*)block_cnts.data_ptr(), + // --- End of Correction --- + nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr + )).wait(); + }); + + at::Tensor block_accum_inclusive = at::cumsum(block_cnts, 0); + + int64_t nnz = 0; // Use int64_t for nnz to match tensor type + 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); + } + + // --- Start of Correction --- + // To get an exclusive scan, shift the inclusive scan and prepend a zero. + // Ensure the types are consistent (kLong). + at::Tensor block_accum_exclusive = at::cat({at::zeros({1}, long_opts), block_accum_inclusive.slice(0, 0, n_blocks - 1)}); + // The kernel expects int32_t*, so we must convert the exclusive scan result back to kInt. + at::Tensor block_accum_exclusive_int = block_accum_exclusive.to(at::kInt); + // --- End of Correction --- + + + // 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, + // --- Start of Correction --- + block_accum_exclusive_int.data_ptr(), + // --- End of Correction --- + 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_({(long)nrows}, nnz); + } + + return std::make_tuple( + batch_ids, camera_ids, gaussian_ids, radii, means2d, depths, conics, indptr, compensations); } } // namespace gsplat::xpu \ No newline at end of file From fc801e821d12120b5d7533eb26d4c9f2d454d97e Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Thu, 25 Sep 2025 09:24:30 +0000 Subject: [PATCH 11/56] added fully fused projection packed --- gsplat/sycl/_wrapper.py | 2 +- .../src/projection_ewa_3dgs_packed_fwd.cpp | 23 ++++++++----------- 2 files changed, 10 insertions(+), 15 deletions(-) diff --git a/gsplat/sycl/_wrapper.py b/gsplat/sycl/_wrapper.py index 2172bf88..92ad3733 100644 --- a/gsplat/sycl/_wrapper.py +++ b/gsplat/sycl/_wrapper.py @@ -1607,7 +1607,6 @@ def forward( ) ( - indptr, batch_ids, camera_ids, gaussian_ids, @@ -1615,6 +1614,7 @@ def forward( means2d, depths, conics, + indptr, compensations, ) = _make_lazy_sycl_func("projection_ewa_3dgs_packed_fwd")( means, diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp index 0c60c812..81402407 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -68,8 +68,8 @@ projection_ewa_3dgs_packed_fwd( } // --- Start of Correction --- - // Changed block_cnts to kLong to satisfy at::cumsum requirements - at::Tensor block_cnts = at::empty({(long)n_blocks}, long_opts); + // Allocate block_cnts as kInt, which the kernel expects. + at::Tensor block_cnts = at::empty({(long)n_blocks}, int_opts); // --- End of Correction --- auto& d_queue = at::xpu::getCurrentXPUStream().queue(); @@ -93,17 +93,19 @@ projection_ewa_3dgs_packed_fwd( camera_model, nullptr, // block_accum // --- Start of Correction --- - // The kernel expects int32_t*, but the tensor is int64_t. - // This cast is safe because the counts per block will not exceed int32_t max. - (int32_t*)block_cnts.data_ptr(), + // Pass the int32_t pointer directly, no cast needed. + block_cnts.data_ptr(), // --- End of Correction --- nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr )).wait(); }); - at::Tensor block_accum_inclusive = at::cumsum(block_cnts, 0); + // --- Start of Correction --- + // 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); + // --- End of Correction --- - int64_t nnz = 0; // Use int64_t for nnz to match tensor type + int64_t nnz = 0; if (n_blocks > 0) { nnz = block_accum_inclusive.index({-1}).item(); } @@ -113,13 +115,8 @@ projection_ewa_3dgs_packed_fwd( batch_ids, camera_ids, gaussian_ids, radii, means2d, depths, conics, indptr, compensations); } - // --- Start of Correction --- - // To get an exclusive scan, shift the inclusive scan and prepend a zero. - // Ensure the types are consistent (kLong). at::Tensor block_accum_exclusive = at::cat({at::zeros({1}, long_opts), block_accum_inclusive.slice(0, 0, n_blocks - 1)}); - // The kernel expects int32_t*, so we must convert the exclusive scan result back to kInt. at::Tensor block_accum_exclusive_int = block_accum_exclusive.to(at::kInt); - // --- End of Correction --- // Allocate final output tensors @@ -148,9 +145,7 @@ projection_ewa_3dgs_packed_fwd( image_width, image_height, (scalar_t)eps2d, (scalar_t)near_plane, (scalar_t)far_plane, (scalar_t)radius_clip, camera_model, - // --- Start of Correction --- block_accum_exclusive_int.data_ptr(), - // --- End of Correction --- nullptr, // block_cnts indptr.data_ptr(), batch_ids.data_ptr(), From f7738591302926134f2a0936304646d13e735b62 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Thu, 25 Sep 2025 09:38:52 +0000 Subject: [PATCH 12/56] Update kernel --- gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp index 23419234..36f4b328 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -162,7 +162,7 @@ struct PackedProjectionFwdKernel { if (opacity < ALPHA_THRESHOLD) { valid = false; } - extend = sycl::min(extend, sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD))); + extend = sycl::fmin(extend, sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD))); } radius_x = sycl::ceil(extend * sycl::sqrt(covar2d[0][0])); @@ -183,7 +183,6 @@ struct PackedProjectionFwdKernel { if (m_block_cnts != nullptr) { // First pass: Count visible Gaussians in this block. - // Check if any thread in the group has a valid Gaussian. bool any_valid = sycl::any_of_group(group, valid); if (any_valid) { // Reduce the count of valid Gaussians across the work-group. From 4dfea003e8a6379112528afe1401648fdcefd22e Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Thu, 25 Sep 2025 14:10:17 +0000 Subject: [PATCH 13/56] updated training code --- examples/requirements.txt | 4 +-- examples/simple_trainer.py | 22 +++++++++++------ gsplat/distributed.py | 45 ++++++++++++++++++++++++++++------ gsplat/rendering.py | 50 ++++++++++++++++++++++++++------------ 4 files changed, 88 insertions(+), 33 deletions(-) diff --git a/examples/requirements.txt b/examples/requirements.txt index ea0a940e..bbf4fede 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/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe +#git+https://github.com/rahul-goel/fused-ssim@328dc9836f513d00c4b5bc38fe30478b4435cbb5 +#git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe splines diff --git a/examples/simple_trainer.py b/examples/simple_trainer.py index 6a30be73..2b02de72 100644 --- a/examples/simple_trainer.py +++ b/examples/simple_trainer.py @@ -21,7 +21,7 @@ generate_interpolated_path, generate_spiral_path, ) -from fused_ssim import fused_ssim +from fusedssim_sycl import fusedssim from torch import Tensor from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.tensorboard import SummaryWriter @@ -227,7 +227,7 @@ def create_splats_with_optimizers( visible_adam: bool = False, batch_size: int = 1, feature_dim: Optional[int] = None, - device: str = "cuda", + device: str = "xpu" if torch.xpu.is_available() else "cuda" , world_rank: int = 0, world_size: int = 1, ) -> Tuple[torch.nn.ParameterDict, Dict[str, torch.optim.Optimizer]]: @@ -312,7 +312,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 = f"xpu:{local_rank}" if torch.xpu.is_available() else f"cuda:{local_rank}" # Where to dump results. os.makedirs(cfg.result_dir, exist_ok=True) @@ -681,7 +681,7 @@ def train(self): # loss l1loss = F.l1_loss(colors, pixels) - ssimloss = 1.0 - fused_ssim( + ssimloss = 1.0 - fusedssim( colors.permute(0, 3, 1, 2), pixels.permute(0, 3, 1, 2), padding="valid" ) loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda @@ -735,7 +735,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.xpu.max_memory_allocated() if torch.xpu.is_available() else torch.cuda.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) @@ -753,7 +753,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.xpu.max_memory_allocated() if torch.xpu.is_available() else torch.cuda.max_memory_allocated()) / 1024**3 stats = { "mem": mem, "ellipse_time": time.time() - global_tic, @@ -923,7 +923,10 @@ 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() + if torch.xpu.is_available(): + torch.xpu.synchronize() + else: + torch.cuda.synchronize() tic = time.time() colors, _, _ = self.rasterize_splats( camtoworlds=camtoworlds, @@ -935,7 +938,10 @@ def eval(self, step: int, stage: str = "val"): far_plane=cfg.far_plane, masks=masks, ) # [1, H, W, 3] - torch.cuda.synchronize() + if torch.xpu.is_available(): + torch.xpu.synchronize() + else: + torch.cuda.synchronize() ellipse_time += max(time.time() - tic, 1e-10) colors = torch.clamp(colors, 0.0, 1.0) diff --git a/gsplat/distributed.py b/gsplat/distributed.py index cab559df..9cf0e66a 100644 --- a/gsplat/distributed.py +++ b/gsplat/distributed.py @@ -6,6 +6,14 @@ import torch.distributed.nn.functional as distF from torch import Tensor +import gsplat + + +def _get_distributed_backend(): + if gsplat.BACKEND == "sycl": + return "ccl" + return "nccl" + def all_gather_int32( world_size: int, value: Union[int, Tensor], device: Optional[torch.device] = None @@ -30,13 +38,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 gsplat.BACKEND == "cuda": + assert value_tensor.is_cuda, "value should be on CUDA" + elif gsplat.BACKEND == "sycl": + assert value_tensor.is_xpu, "value should be on XPU" # gather collected = torch.empty( @@ -82,7 +94,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 +295,15 @@ 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) + if gsplat.BACKEND == "cuda": + torch.cuda.set_device(local_rank) + elif gsplat.BACKEND == "sycl": + import torch_ccl + import torch.xpu + torch.xpu.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 +337,12 @@ 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 gsplat.BACKEND == "cuda": + assert torch.cuda.is_available(), "CUDA device is required!" + elif gsplat.BACKEND == "sycl": + import torch.xpu + assert torch.xpu.is_available(), "XPU 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 +351,13 @@ 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() + if gsplat.BACKEND == "cuda": + world_size = torch.cuda.device_count() + elif gsplat.BACKEND == "sycl": + world_size = torch.xpu.device_count() + else: + world_size = 1 + distributed = world_size > 1 if distributed: @@ -357,4 +386,4 @@ def fn(local_rank: int, world_rank: int, world_size: int, args: Any) -> None: print("process " + str(i) + " finished") return True else: - return _distributed_worker(0, 1, fn=fn, args=args) + return _distributed_worker(0, 1, fn=fn, args=args) \ No newline at end of file diff --git a/gsplat/rendering.py b/gsplat/rendering.py index 2054fed2..cee97e2a 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -7,20 +7,40 @@ from torch import Tensor from typing_extensions import Literal -from .cuda._wrapper import ( - RollingShutterType, - FThetaCameraDistortionParameters, - FThetaPolynomialType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, -) +from . import BACKEND + +# Now, conditionally import the functions based on the detected backend. +if BACKEND == "cuda": + from .cuda._wrapper import ( + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + ) +elif BACKEND == "sycl": + from .sycl._wrapper import ( + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + ) +else: + # If no backend is found, you can either raise an error or define dummy functions + # to avoid crashing, depending on your needs. + raise ImportError("gsplat: No backend (CUDA or SYCL) found, cannot import backend-specific functions.") + from .distributed import ( all_gather_int32, all_gather_tensor_list, @@ -63,7 +83,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] From 6624ec0a2137fadb216bc054df922204db8002b1 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Thu, 25 Sep 2025 14:18:30 +0000 Subject: [PATCH 14/56] update steps for memory error --- examples/simple_trainer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/simple_trainer.py b/examples/simple_trainer.py index 2b02de72..057656d0 100644 --- a/examples/simple_trainer.py +++ b/examples/simple_trainer.py @@ -79,13 +79,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 From 6597d44ac5b078d2b835cdf46e4cebf2e447de22 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Tue, 7 Oct 2025 15:32:45 -0700 Subject: [PATCH 15/56] torch_acc uniform API for both cuda and xpu. Maybe replaced by torch.accelerator in the future. Ensure empty_cache() and synchronize() for XPU match those for CUDA, --- examples/image_fitting.py | 8 +++--- examples/simple_trainer.py | 51 ++++++++++++++++++++++++--------- examples/simple_trainer_2dgs.py | 13 +++++---- gsplat/__init__.py | 3 ++ gsplat/distributed.py | 15 ++-------- gsplat/profile.py | 5 ++-- gsplat/strategy/default.py | 5 ++-- gsplat/strategy/mcmc.py | 3 +- profiling/batch.py | 19 ++++++------ profiling/main.py | 25 +++++++++------- 10 files changed, 88 insertions(+), 59 deletions(-) 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/simple_trainer.py b/examples/simple_trainer.py index 057656d0..9f14bd25 100644 --- a/examples/simple_trainer.py +++ b/examples/simple_trainer.py @@ -30,7 +30,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 from gsplat.compression import PngCompression from gsplat.distributed import cli from gsplat.optimizers import SelectiveAdam @@ -227,7 +227,7 @@ def create_splats_with_optimizers( visible_adam: bool = False, batch_size: int = 1, feature_dim: Optional[int] = None, - device: str = "xpu" if torch.xpu.is_available() else "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]]: @@ -312,7 +312,7 @@ def __init__( self.world_rank = world_rank self.local_rank = local_rank self.world_size = world_size - self.device = f"xpu:{local_rank}" if torch.xpu.is_available() else f"cuda:{local_rank}" + self.device = str(torch_acc._device(local_rank)) # Where to dump results. os.makedirs(cfg.result_dir, exist_ok=True) @@ -735,7 +735,7 @@ def train(self): # ) if world_rank == 0 and cfg.tb_every > 0 and step % cfg.tb_every == 0: - mem = (torch.xpu.max_memory_allocated() if torch.xpu.is_available() else 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) @@ -753,7 +753,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.xpu.max_memory_allocated() if torch.xpu.is_available() else torch.cuda.max_memory_allocated()) / 1024**3 + mem = torch_acc.max_memory_allocated() / 1024**3 stats = { "mem": mem, "ellipse_time": time.time() - global_tic, @@ -923,10 +923,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] - if torch.xpu.is_available(): - torch.xpu.synchronize() - else: - torch.cuda.synchronize() + torch_acc.synchronize() tic = time.time() colors, _, _ = self.rasterize_splats( camtoworlds=camtoworlds, @@ -938,10 +935,7 @@ def eval(self, step: int, stage: str = "val"): far_plane=cfg.far_plane, masks=masks, ) # [1, H, W, 3] - if torch.xpu.is_available(): - torch.xpu.synchronize() - else: - torch.cuda.synchronize() + torch_acc.synchronize() ellipse_time += max(time.time() - tic, 1e-10) colors = torch.clamp(colors, 0.0, 1.0) @@ -1183,6 +1177,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/gsplat/__init__.py b/gsplat/__init__.py index 5cf505b0..14449eb3 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -27,6 +27,7 @@ spherical_harmonics, world_to_cam, ) + torch_acc = torch.cuda print("gsplat: CUDA backend successfully loaded.", file=sys.stderr) except ImportError: if FORCE_BACKEND == "cuda": @@ -53,6 +54,7 @@ spherical_harmonics, world_to_cam, ) + torch_acc = torch.xpu print("gsplat: SYCL backend successfully loaded.", file=sys.stderr) except ImportError as e: if FORCE_BACKEND == "sycl": @@ -81,6 +83,7 @@ __all__ = [ "BACKEND", + "torch_acc", "PngCompression", "DefaultStrategy", "MCMCStrategy", diff --git a/gsplat/distributed.py b/gsplat/distributed.py index 9cf0e66a..26e2de23 100644 --- a/gsplat/distributed.py +++ b/gsplat/distributed.py @@ -7,6 +7,7 @@ from torch import Tensor import gsplat +from gsplat import torch_acc def _get_distributed_backend(): @@ -295,12 +296,7 @@ def _distributed_worker( print("Distributed worker: %d / %d" % (world_rank + 1, world_size)) distributed = world_size > 1 if distributed: - if gsplat.BACKEND == "cuda": - torch.cuda.set_device(local_rank) - elif gsplat.BACKEND == "sycl": - import torch_ccl - import torch.xpu - torch.xpu.set_device(local_rank) + torch_acc.set_device(local_rank) torch.distributed.init_process_group( backend=_get_distributed_backend(), world_size=world_size, rank=world_rank @@ -351,13 +347,8 @@ def fn(local_rank: int, world_rank: int, world_size: int, args: Any) -> None: world_rank, world_size, fn, args, local_rank, verbose ) - if gsplat.BACKEND == "cuda": - world_size = torch.cuda.device_count() - elif gsplat.BACKEND == "sycl": - world_size = torch.xpu.device_count() - else: - world_size = 1 + world_size = torch_acc.device_count() distributed = world_size > 1 if distributed: 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/strategy/default.py b/gsplat/strategy/default.py index 49e677d9..61f68aee 100644 --- a/gsplat/strategy/default.py +++ b/gsplat/strategy/default.py @@ -6,7 +6,7 @@ from .base import Strategy from .ops import duplicate, remove, reset_opa, split - +from .. import torch_acc @dataclass class DefaultStrategy(Strategy): @@ -190,7 +190,8 @@ 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() + print(f"Empty cache after step {step}", flush=True) if step % self.reset_every == 0 & 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/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..1e156d84 100644 --- a/profiling/main.py +++ b/profiling/main.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 @@ -86,8 +87,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 +121,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 +132,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 " @@ -194,7 +195,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: @@ -225,7 +226,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: @@ -256,7 +257,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 +286,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 +367,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) From 0a472cfeb642753599397b6425596f6ab4938740 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Fri, 17 Oct 2025 13:47:14 +0000 Subject: [PATCH 16/56] Added forward pass for 2fgs fully fused projection --- .../kernels/Projection2DGSFusedFwdKernel.hpp | 212 ++++++++++++++++++ gsplat/sycl/src/projection_2dgs_fused_fwd.cpp | 89 +++++++- 2 files changed, 297 insertions(+), 4 deletions(-) create mode 100644 gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp new file mode 100644 index 00000000..5598dc5f --- /dev/null +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp @@ -0,0 +1,212 @@ +#ifndef Projection2DGSFusedFwdKernel_HPP +#define Projection2DGSFusedFwdKernel_HPP + +#include "utils.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.hpp" + +namespace gsplat::xpu { + +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][0], RS_camera[1][0], mean_c.x, + RS_camera[0][1], RS_camera[1][1], mean_c.y, + RS_camera[0][2], RS_camera[1][2], mean_c.z + ); + + // 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 = dot(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( + dot(f * M0, M2), + dot(f * M1, M2) + ); + + const vec2 temp = vec2( + dot(f * M0, M0), + dot(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]); + + // 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] = M0.x; + m_ray_transforms[idx * 9 + 1] = M0.y; + m_ray_transforms[idx * 9 + 2] = M0.z; + m_ray_transforms[idx * 9 + 3] = M1.x; + m_ray_transforms[idx * 9 + 4] = M1.y; + m_ray_transforms[idx * 9 + 5] = M1.z; + m_ray_transforms[idx * 9 + 6] = M2.x; + m_ray_transforms[idx * 9 + 7] = M2.y; + m_ray_transforms[idx * 9 + 8] = M2.z; + + // 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 + +#endif // Projection2DGSFusedFwdKernel_HPP \ 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 index c5167a00..dec9a0d5 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp @@ -1,10 +1,10 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/Projection2DGSFusedFwdKernel.hpp" -namespace gsplat::xpu { +namespace gsplat::xpu { std::tuple< at::Tensor, @@ -25,7 +25,88 @@ projection_2dgs_fused_fwd( const float far_plane, const float radius_clip ) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + CHECK_CONTIGUOUS(means); + CHECK_CONTIGUOUS(quats); + CHECK_CONTIGUOUS(scales); + CHECK_CONTIGUOUS(viewmats); + CHECK_CONTIGUOUS(Ks); + + 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] (flattened to [..., C, N, 9]) + at::DimVector out_shape_cn33 = batch_dims; + out_shape_cn33.insert(out_shape_cn33.end(), {C, N, 9}); + + 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 +} // namespace gsplat::xpu \ No newline at end of file From 417942e90ae394e54c0d4ab23b000b9be7ac3d2e Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Fri, 17 Oct 2025 14:33:12 +0000 Subject: [PATCH 17/56] added backward pass --- .../kernels/Projection2DGSFusedBwdKernel.hpp | 279 ++++++++++++++++++ gsplat/sycl/src/projection_2dgs_fused_bwd.cpp | 79 ++++- 2 files changed, 354 insertions(+), 4 deletions(-) create mode 100644 gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp new file mode 100644 index 00000000..7272e6ae --- /dev/null +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp @@ -0,0 +1,279 @@ +#ifndef Projection2DGSFusedBwdKernel_HPP +#define Projection2DGSFusedBwdKernel_HPP + +#include "utils.hpp" +#include "quat_scale_to_covar_preci.hpp" +#include "transform.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 + +#endif // Projection2DGSFusedBwdKernel_HPP \ 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 index 945ad018..2b08b134 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp @@ -1,10 +1,10 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/Projection2DGSFusedBwdKernel.hpp" -namespace gsplat::xpu { +namespace gsplat::xpu { std::tuple projection_2dgs_fused_bwd( @@ -26,7 +26,78 @@ projection_2dgs_fused_bwd( const at::Tensor v_ray_transforms, // [..., C, N, 3, 3] const bool viewmats_requires_grad ) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + CHECK_CONTIGUOUS(means); + CHECK_CONTIGUOUS(quats); + CHECK_CONTIGUOUS(scales); + CHECK_CONTIGUOUS(viewmats); + CHECK_CONTIGUOUS(Ks); + CHECK_CONTIGUOUS(radii); + CHECK_CONTIGUOUS(ray_transforms); + CHECK_CONTIGUOUS(v_means2d); + CHECK_CONTIGUOUS(v_depths); + CHECK_CONTIGUOUS(v_normals); + CHECK_CONTIGUOUS(v_ray_transforms); + + 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 +} // namespace gsplat::xpu \ No newline at end of file From 6763adac2995608f7964f8aff867de4dd4448146 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Fri, 17 Oct 2025 14:45:46 +0000 Subject: [PATCH 18/56] update tests --- tests/test_2dgs.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 788c9111..a78efa32 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -4,7 +4,20 @@ 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 = None + +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="Test requires CUDA backend" +) def expand(data: dict, batch_dims: Tuple[int, ...]): @@ -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) From f4d0123e068f7856e37c18c3a02befdf3e77fbfa Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Fri, 17 Oct 2025 15:28:33 +0000 Subject: [PATCH 19/56] Working tests --- .../kernels/Projection2DGSFusedFwdKernel.hpp | 43 +++++++++---------- gsplat/sycl/src/projection_2dgs_fused_fwd.cpp | 4 +- tests/test_2dgs.py | 1 + 3 files changed, 23 insertions(+), 25 deletions(-) diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp index 5598dc5f..8afa122f 100644 --- a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp @@ -7,6 +7,9 @@ namespace gsplat::xpu { +template +inline float sum(vec3 a) { return a.x + a.y + a.z; } + template struct Projection2DGSFusedFwdKernel { const uint32_t m_B; @@ -113,9 +116,7 @@ struct Projection2DGSFusedFwdKernel { // WH = [RS_camera[0], RS_camera[1], mean_c] mat3 WH = mat3( - RS_camera[0][0], RS_camera[1][0], mean_c.x, - RS_camera[0][1], RS_camera[1][1], mean_c.y, - RS_camera[0][2], RS_camera[1][2], mean_c.z + RS_camera[0], RS_camera[1], mean_c ); // Projective transformation matrix: Camera -> Screen @@ -137,7 +138,7 @@ struct Projection2DGSFusedFwdKernel { const vec3 temp_point = vec3(T(1.0), T(1.0), T(-1.0)); // Algebraic manipulation for computing mean and radius - const T distance = dot(temp_point * M2, M2); + const T distance = sum(temp_point * M2 * M2); // Ignore ill-conditioned primitives if (distance == T(0.0)) { @@ -147,15 +148,9 @@ struct Projection2DGSFusedFwdKernel { } const vec3 f = (T(1.0) / distance) * temp_point; - const vec2 mean2d = vec2( - dot(f * M0, M2), - dot(f * M1, M2) - ); + 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 temp = vec2( - dot(f * M0, M0), - dot(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))); @@ -176,8 +171,9 @@ struct Projection2DGSFusedFwdKernel { } // Compute normals (dual visible) - vec3 normal = vec3(RS_camera[2][0], RS_camera[2][1], RS_camera[2][2]); - + // 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; @@ -190,15 +186,16 @@ struct Projection2DGSFusedFwdKernel { m_depths[idx] = mean_c.z; // Store ray transforms (row major KWH) - m_ray_transforms[idx * 9] = M0.x; - m_ray_transforms[idx * 9 + 1] = M0.y; - m_ray_transforms[idx * 9 + 2] = M0.z; - m_ray_transforms[idx * 9 + 3] = M1.x; - m_ray_transforms[idx * 9 + 4] = M1.y; - m_ray_transforms[idx * 9 + 5] = M1.z; - m_ray_transforms[idx * 9 + 6] = M2.x; - m_ray_transforms[idx * 9 + 7] = M2.y; - m_ray_transforms[idx * 9 + 8] = M2.z; + 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; diff --git a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp index dec9a0d5..731e218b 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp @@ -57,9 +57,9 @@ projection_2dgs_fused_fwd( at::DimVector out_shape_cn3 = batch_dims; out_shape_cn3.insert(out_shape_cn3.end(), {C, N, 3}); - // Output shape: [..., C, N, 3, 3] (flattened to [..., C, N, 9]) + // Output shape: [..., C, N, 3, 3] at::DimVector out_shape_cn33 = batch_dims; - out_shape_cn33.insert(out_shape_cn33.end(), {C, N, 9}); + 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); diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index a78efa32..fd38d1f5 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -100,6 +100,7 @@ def test_projection_2dgs(test_data, batch_dims: Tuple[int, ...]): # TODO (WZ): is the following true for 2dgs as while? # radii is integer so we allow for 1 unit difference valid = ((radii > 0) & (_radii > 0)).all(dim=-1) + valid_expanded = valid.unsqueeze(-1).unsqueeze(-1) torch.testing.assert_close(radii, _radii, rtol=1e-3, atol=1) torch.testing.assert_close(means2d[valid], _means2d[valid], rtol=1e-4, atol=1e-4) torch.testing.assert_close(depths[valid], _depths[valid], rtol=1e-4, atol=1e-4) From 167d0b6fbb70b368b1157d8f4b7e5b0484072773 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Fri, 17 Oct 2025 15:42:10 +0000 Subject: [PATCH 20/56] Added rasterize forward kernel --- .../RasterizeToPixels2DGSFwdKernel.hpp | 392 ++++++++++++++++++ .../sycl/src/rasterize_to_pixels_2dgs_fwd.cpp | 165 +++++++- 2 files changed, 552 insertions(+), 5 deletions(-) create mode 100644 gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp new file mode 100644 index 00000000..796924a9 --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp @@ -0,0 +1,392 @@ +#ifndef RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP +#define RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP + +#include +#include "types.hpp" +#include "gsplat_sycl_utils.hpp" + +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 (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*>(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 + +#endif // RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp index 4a479d11..3a7dc23b 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp @@ -1,11 +1,96 @@ - #include #include "Ops.h" #include "Common.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_width}; + 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 -namespace gsplat::xpu { - std::tuple< at::Tensor, at::Tensor, @@ -31,7 +116,77 @@ rasterize_to_pixels_2dgs_fwd( 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"); + // Check input tensors are contiguous + CHECK_CONTIGUOUS(means2d); + CHECK_CONTIGUOUS(ray_transforms); + CHECK_CONTIGUOUS(colors); + CHECK_CONTIGUOUS(opacities); + CHECK_CONTIGUOUS(normals); + CHECK_CONTIGUOUS(tile_offsets); + CHECK_CONTIGUOUS(flatten_ids); + if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + + // Get dimensions + bool packed = means2d.dim() == 2; + uint32_t N = packed ? 0 : means2d.size(-2); // number of gaussians + uint32_t I = tile_offsets.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); // 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::Tensor renders = at::zeros({I, image_height, image_width, channels}, options_float); + at::Tensor alphas = at::zeros({I, image_height, image_width}, options_float); + at::Tensor render_normals = at::zeros({I, image_height, image_width, 3}, options_float); + at::Tensor render_distort = at::zeros({I, image_height, image_width}, options_float); + at::Tensor render_median = at::zeros({I, image_height, image_width}, options_float); + at::Tensor last_ids = at::zeros({I, image_height, image_width}, options_int); + at::Tensor median_ids = at::zeros({I, image_height, image_width}, 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 \ No newline at end of file +} // namespace gsplat::xpu \ No newline at end of file From a4f90ac698daaf6d191a9f0162df369269c8153c Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Fri, 17 Oct 2025 16:07:06 +0000 Subject: [PATCH 21/56] added backward kernel --- .../RasterizeToPixels2DGSBwdKernel.hpp | 614 ++++++++++++++++++ .../sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 259 +++++++- tests/test_2dgs.py | 2 +- 3 files changed, 847 insertions(+), 28 deletions(-) create mode 100644 gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp new file mode 100644 index 00000000..29e14a98 --- /dev/null +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp @@ -0,0 +1,614 @@ +#ifndef RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP +#define RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP + +#include +#include "types.hpp" +#include "gsplat_sycl_utils.hpp" + +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 float* m_render_normals; // Rendered normals + const float* m_render_distort; // Distortion values + const float* m_render_median; // Median depth 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 float* render_normals, + const float* render_distort, + const float* render_median, + 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_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_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 float* render_normals_ptr = m_render_normals + image_id * m_image_height * m_image_width * 3; + const float* render_distort_ptr = nullptr; + if (m_render_distort != nullptr) { + render_distort_ptr = m_render_distort + image_id * m_image_height * m_image_width; + } + const float* render_median_ptr = m_render_median + image_id * m_image_height * m_image_width; + + 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 = num_batches - 1; b >= 0; --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*>(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 + +#endif // RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP \ 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 index 2ad95e89..baba8a9c 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -1,51 +1,256 @@ - #include #include "Ops.h" #include "Common.h" +#include "kernels/RasterizeToPixels2DGSBwdKernel.hpp" -namespace gsplat::xpu { +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& render_normals, + const at::Tensor& render_distort, + const at::Tensor& render_median, + 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_width}; + 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(), + render_normals.data_ptr(), + render_distort.data_ptr(), + render_median.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, 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] + 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, - // ray_crossions - const at::Tensor tile_offsets, // [..., tile_height, tile_width] - const at::Tensor flatten_ids, // [n_isects] + // 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, 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] + const at::Tensor render_colors, // [..., image_height, image_width, channels] + const at::Tensor render_alphas, // [..., image_height, image_width] + const at::Tensor render_normals, // [..., image_height, image_width, 3] + const at::Tensor render_distort, // [..., image_height, image_width] + const at::Tensor render_median, // [..., 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, 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 + 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] ) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + // Check input tensors are contiguous + CHECK_CONTIGUOUS(means2d); + CHECK_CONTIGUOUS(ray_transforms); + CHECK_CONTIGUOUS(colors); + CHECK_CONTIGUOUS(opacities); + CHECK_CONTIGUOUS(normals); + CHECK_CONTIGUOUS(densify); + CHECK_CONTIGUOUS(tile_offsets); + CHECK_CONTIGUOUS(flatten_ids); + CHECK_CONTIGUOUS(render_colors); + CHECK_CONTIGUOUS(render_alphas); + CHECK_CONTIGUOUS(render_normals); + CHECK_CONTIGUOUS(render_distort); + CHECK_CONTIGUOUS(render_median); + CHECK_CONTIGUOUS(last_ids); + CHECK_CONTIGUOUS(median_ids); + CHECK_CONTIGUOUS(v_render_colors); + CHECK_CONTIGUOUS(v_render_alphas); + CHECK_CONTIGUOUS(v_render_normals); + CHECK_CONTIGUOUS(v_render_distort); + CHECK_CONTIGUOUS(v_render_median); + if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + + uint32_t channels = colors.size(-1); + bool compute_mean_abs = true; + + // Create output tensors + auto options = means2d.options().dtype(torch::kFloat32); + at::Tensor v_means2d_abs = compute_mean_abs ? + at::zeros_like(means2d, options) : at::Tensor(); + 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, \ + render_normals, render_distort, render_median, last_ids, median_ids, \ + v_render_colors, v_render_alphas, v_render_normals, \ + v_render_distort, v_render_median, \ + compute_mean_abs ? 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( + compute_mean_abs ? v_means2d_abs : v_means2d, + v_means2d, + v_ray_transforms, + v_colors, + v_opacities, + v_normals + ); } -} // namespace gsplat::xpu \ No newline at end of file +} // namespace gsplat::xpu \ No newline at end of file diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index fd38d1f5..58752551 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -268,7 +268,7 @@ def test_fully_fused_projection_packed_2dgs( 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._torch_impl_2dgs import _rasterize_to_pixels_2dgs from gsplat.cuda._wrapper import ( fully_fused_projection_2dgs, isect_offset_encode, From aa6ff0809d539b8ac57e46ebc3462128115c175a Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Tue, 11 Nov 2025 15:50:07 +0530 Subject: [PATCH 22/56] Update rasterize_to_pixels_2dgs_bwd.cpp --- .../sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 52 ++++++++----------- 1 file changed, 22 insertions(+), 30 deletions(-) diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp index baba8a9c..71012bd9 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -29,9 +29,6 @@ void launch_rasterize_2dgs_bwd_kernel( // forward outputs const at::Tensor& render_colors, const at::Tensor& render_alphas, - const at::Tensor& render_normals, - const at::Tensor& render_distort, - const at::Tensor& render_median, const at::Tensor& last_ids, const at::Tensor& median_ids, // gradients of outputs @@ -97,9 +94,6 @@ void launch_rasterize_2dgs_bwd_kernel( flatten_ids.data_ptr(), render_colors.data_ptr(), render_alphas.data_ptr(), - render_normals.data_ptr(), - render_distort.data_ptr(), - render_median.data_ptr(), last_ids.data_ptr(), median_ids.data_ptr(), v_render_colors.data_ptr(), @@ -132,6 +126,7 @@ std::tuple< at::Tensor, at::Tensor, at::Tensor, + at::Tensor, at::Tensor> rasterize_to_pixels_2dgs_bwd( // Gaussian parameters @@ -153,9 +148,6 @@ rasterize_to_pixels_2dgs_bwd( // forward outputs const at::Tensor render_colors, // [..., image_height, image_width, channels] const at::Tensor render_alphas, // [..., image_height, image_width] - const at::Tensor render_normals, // [..., image_height, image_width, 3] - const at::Tensor render_distort, // [..., image_height, image_width] - const at::Tensor render_median, // [..., 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 @@ -163,7 +155,8 @@ rasterize_to_pixels_2dgs_bwd( 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] + const at::Tensor v_render_median, // [..., image_height, image_width] + bool absgrad ) { // Check input tensors are contiguous CHECK_CONTIGUOUS(means2d); @@ -176,9 +169,6 @@ rasterize_to_pixels_2dgs_bwd( CHECK_CONTIGUOUS(flatten_ids); CHECK_CONTIGUOUS(render_colors); CHECK_CONTIGUOUS(render_alphas); - CHECK_CONTIGUOUS(render_normals); - CHECK_CONTIGUOUS(render_distort); - CHECK_CONTIGUOUS(render_median); CHECK_CONTIGUOUS(last_ids); CHECK_CONTIGUOUS(median_ids); CHECK_CONTIGUOUS(v_render_colors); @@ -190,12 +180,13 @@ rasterize_to_pixels_2dgs_bwd( if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); uint32_t channels = colors.size(-1); - bool compute_mean_abs = true; // Create output tensors auto options = means2d.options().dtype(torch::kFloat32); - at::Tensor v_means2d_abs = compute_mean_abs ? - at::zeros_like(means2d, options) : at::Tensor(); + 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); @@ -204,18 +195,18 @@ rasterize_to_pixels_2dgs_bwd( 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, \ - render_normals, render_distort, render_median, last_ids, median_ids, \ - v_render_colors, v_render_alphas, v_render_normals, \ - v_render_distort, v_render_median, \ - compute_mean_abs ? c10::optional(v_means2d_abs) : c10::nullopt, \ +#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) { @@ -244,13 +235,14 @@ rasterize_to_pixels_2dgs_bwd( #undef __GS__CALL_ return std::make_tuple( - compute_mean_abs ? v_means2d_abs : v_means2d, + v_means2d_abs, v_means2d, v_ray_transforms, v_colors, v_opacities, - v_normals + v_normals, + v_densify ); } -} // namespace gsplat::xpu \ No newline at end of file +} // namespace gsplat::xpu From 6542d4c4892a9b3c47fe20ad81a6c2406acfe554 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Tue, 11 Nov 2025 15:56:26 +0530 Subject: [PATCH 23/56] Update RasterizeToPixels2DGSBwdKernel.hpp --- .../kernels/RasterizeToPixels2DGSBwdKernel.hpp | 17 ++--------------- 1 file changed, 2 insertions(+), 15 deletions(-) diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp index 29e14a98..b939e4fe 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp @@ -43,9 +43,6 @@ struct RasterizeToPixels2DGSBwdKernel { // Forward pass outputs const float* m_render_colors; // Rendered colors const float* m_render_alphas; // Alpha values - const float* m_render_normals; // Rendered normals - const float* m_render_distort; // Distortion values - const float* m_render_median; // Median depth values const int32_t* m_last_ids; // Last Gaussian indices const int32_t* m_median_ids; // Median Gaussian indices @@ -100,9 +97,6 @@ struct RasterizeToPixels2DGSBwdKernel { // Forward outputs const float* render_colors, const float* render_alphas, - const float* render_normals, - const float* render_distort, - const float* render_median, const int32_t* last_ids, const int32_t* median_ids, // Gradient inputs @@ -136,8 +130,7 @@ struct RasterizeToPixels2DGSBwdKernel { 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_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), @@ -163,12 +156,6 @@ struct RasterizeToPixels2DGSBwdKernel { 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 float* render_normals_ptr = m_render_normals + image_id * m_image_height * m_image_width * 3; - const float* render_distort_ptr = nullptr; - if (m_render_distort != nullptr) { - render_distort_ptr = m_render_distort + image_id * m_image_height * m_image_width; - } - const float* render_median_ptr = m_render_median + image_id * m_image_height * m_image_width; 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; @@ -611,4 +598,4 @@ struct RasterizeToPixels2DGSBwdKernel { } // namespace gsplat::xpu -#endif // RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP \ No newline at end of file +#endif // RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP From cb5c72a3b07e4a9d287f869829d7250b97834e41 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Tue, 11 Nov 2025 16:28:15 +0530 Subject: [PATCH 24/56] Update RasterizeToPixels2DGSBwdKernel.hpp --- gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp index b939e4fe..6ceffb12 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp @@ -267,7 +267,7 @@ struct RasterizeToPixels2DGSBwdKernel { ); // Process batches of gaussians in reverse order (back to front) - for (int32_t b = num_batches - 1; b >= 0; --b) { + for (int32_t b = 0; b < num_batches; ++b) { // Synchronize threads before loading next batch item.barrier(sycl::access::fence_space::local_space); From 4a5bde31e59a836ff24abe527e8b69e454211960 Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Tue, 11 Nov 2025 16:54:54 +0530 Subject: [PATCH 25/56] Update RasterizeToPixels2DGSFwdKernel.hpp --- .../sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp index 796924a9..6bcd5b56 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp @@ -183,7 +183,7 @@ struct RasterizeToPixels2DGSFwdKernel { uint32_t batch_start = range_start + m_chunk_size * b; uint32_t idx = batch_start + tr; - if (idx < range_end) { + if (tr < m_chunk_size && idx < range_end) { // Get gaussian index int32_t g = m_flatten_ids[idx]; m_slm_id_batch[tr] = g; @@ -389,4 +389,4 @@ struct RasterizeToPixels2DGSFwdKernel { } // namespace gsplat::xpu -#endif // RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP \ No newline at end of file +#endif // RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP From 11604f542996d35ae4cce711cdba1ab008e8e70a Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Sat, 15 Nov 2025 18:49:29 +0530 Subject: [PATCH 26/56] Update rasterize_to_pixels_2dgs_fwd.cpp for correct block size computation --- gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp index 3a7dc23b..49fed8db 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp @@ -45,7 +45,7 @@ void launch_rasterize_2dgs_kernel( // Define the execution ranges sycl::range<3> localRange{1, tile_size, tile_size}; - sycl::range<3> globalRange{I, tile_height, tile_width}; + 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 @@ -189,4 +189,4 @@ rasterize_to_pixels_2dgs_fwd( render_median, last_ids, median_ids); } -} // namespace gsplat::xpu \ No newline at end of file +} // namespace gsplat::xpu From b5c46a9d8cdc4f1fea27c6497d2739023f23153d Mon Sep 17 00:00:00 2001 From: Aditya Singh Rathore Date: Sat, 15 Nov 2025 18:50:03 +0530 Subject: [PATCH 27/56] Update rasterize_to_pixels_2dgs_bwd.cpp for correct block size computation --- gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp index 71012bd9..8120ce99 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -62,7 +62,7 @@ void launch_rasterize_2dgs_bwd_kernel( // Define the execution ranges sycl::range<3> localRange{1, tile_size, tile_size}; - sycl::range<3> globalRange{I, tile_height, tile_width}; + 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 From eb2174c946d7990c6eb84fbb5740efa08f8f2b3f Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Sat, 15 Nov 2025 08:25:54 -0800 Subject: [PATCH 28/56] relocation kernel for MCMC strategy test updates to always run sycll code (not test if nerfacc not available) rasterize_to_pixels work group size fix, batch dim fix. --- gsplat/__init__.py | 22 +- gsplat/relocation.py | 10 +- gsplat/rendering.py | 2 +- gsplat/sycl/_wrapper.py | 2 +- .../sycl/include/kernels/RelocationKernel.hpp | 57 ++++ .../src/projection_ewa_3dgs_fused_fwd.cpp | 64 +++-- .../sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 8 +- .../sycl/src/rasterize_to_pixels_2dgs_fwd.cpp | 247 +++++++++++------- .../sycl/src/rasterize_to_pixels_3dgs_bwd.cpp | 25 +- .../sycl/src/rasterize_to_pixels_3dgs_fwd.cpp | 34 ++- gsplat/sycl/src/relocation.cpp | 40 ++- tests/test_2dgs.py | 95 ++++--- tests/test_basic.py | 54 ++-- tests/test_rasterization.py | 45 ++-- tests/test_strategy.py | 22 +- 15 files changed, 480 insertions(+), 247 deletions(-) create mode 100644 gsplat/sycl/include/kernels/RelocationKernel.hpp diff --git a/gsplat/__init__.py b/gsplat/__init__.py index 14449eb3..527c77ad 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -27,14 +27,18 @@ spherical_harmonics, world_to_cam, ) + torch_acc = torch.cuda print("gsplat: CUDA backend successfully loaded.", file=sys.stderr) except ImportError: if FORCE_BACKEND == "cuda": - print("gsplat: Error! GSPLAT_BACKEND=cuda was set but CUDA backend failed to load.", file=sys.stderr) + print( + "gsplat: Error! GSPLAT_BACKEND=cuda was set but CUDA backend failed to load.", + file=sys.stderr, + ) pass -if not BACKEND and (FORCE_BACKEND == "sycl" or FORCE_BACKEND == ""): +if not BACKEND and (FORCE_BACKEND in ("sycl", "xpu") or FORCE_BACKEND == "" and torch.xpu.is_available()): try: BACKEND = "sycl" from .sycl._wrapper import ( @@ -54,16 +58,20 @@ spherical_harmonics, world_to_cam, ) + torch_acc = torch.xpu - print("gsplat: SYCL backend successfully loaded.", file=sys.stderr) + print("gsplat: SYCL XPU backend successfully loaded.", file=sys.stderr) except ImportError as e: - if FORCE_BACKEND == "sycl": - print(f"gsplat: Error! GSPLAT_BACKEND=sycl was set but SYCL backend failed to load: {e}", file=sys.stderr) + if FORCE_BACKEND in ("sycl", "xpu"): + print( + f"gsplat: Error! GSPLAT_BACKEND={FORCE_BACKEND} was set but SYCL XPU backend failed to load: {e}", + file=sys.stderr, + ) pass if not BACKEND: print( - "gsplat: Warning! No high-performance backend (CUDA or SYCL) found.", + "gsplat: Warning! No high-performance backend (CUDA or SYCL XPU) found.", file=sys.stderr, ) @@ -111,4 +119,4 @@ "__version__", "SelectiveAdam", # Note: accumulate and accumulate_2dgs are not typically part of the public API -] \ No newline at end of file +] diff --git a/gsplat/relocation.py b/gsplat/relocation.py index 8abd9aee..92b6ce93 100644 --- a/gsplat/relocation.py +++ b/gsplat/relocation.py @@ -4,7 +4,13 @@ import torch from torch import Tensor -from .cuda._wrapper import _make_lazy_cuda_func +from . import BACKEND + +# Now, conditionally import the functions based on the detected backend. +if BACKEND == "cuda": + from .cuda._wrapper import _make_lazy_cuda_func as _make_lazy_func +elif BACKEND == "sycl": + from .sycl._wrapper import _make_lazy_sycl_func as _make_lazy_func def compute_relocation( @@ -43,7 +49,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_func("relocation")( opacities, scales, ratios, binoms, n_max ) return new_opacities, new_scales diff --git a/gsplat/rendering.py b/gsplat/rendering.py index cee97e2a..3e88be4e 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -1498,7 +1498,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/sycl/_wrapper.py b/gsplat/sycl/_wrapper.py index 92ad3733..0f32cb01 100644 --- a/gsplat/sycl/_wrapper.py +++ b/gsplat/sycl/_wrapper.py @@ -2576,7 +2576,7 @@ def backward( v_render_median.contiguous(), absgrad, ) - torch.sycl.synchronize() + torch.xpu.synchronize() if absgrad: means2d.absgrad = v_means2d_abs diff --git a/gsplat/sycl/include/kernels/RelocationKernel.hpp b/gsplat/sycl/include/kernels/RelocationKernel.hpp new file mode 100644 index 00000000..29c544de --- /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/src/projection_ewa_3dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp index 1e42afee..3f9c9f01 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp @@ -1,10 +1,10 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/FullyFusedProjectionFwdKernel.hpp" -namespace gsplat::xpu { +namespace gsplat::xpu { std::tuple< at::Tensor, @@ -32,17 +32,26 @@ projection_ewa_3dgs_fused_fwd( CHECK_CONTIGUOUS(means); CHECK_CONTIGUOUS(viewmats); CHECK_CONTIGUOUS(Ks); - if (covars.has_value()) CHECK_CONTIGUOUS(covars.value()); - if (quats.has_value()) CHECK_CONTIGUOUS(quats.value()); - if (scales.has_value()) CHECK_CONTIGUOUS(scales.value()); - if (opacities.has_value()) CHECK_CONTIGUOUS(opacities.value()); + if (covars.has_value()) + CHECK_CONTIGUOUS(covars.value()); + if (quats.has_value()) + CHECK_CONTIGUOUS(quats.value()); + if (scales.has_value()) + CHECK_CONTIGUOUS(scales.value()); + if (opacities.has_value()) + CHECK_CONTIGUOUS(opacities.value()); - 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]"); + 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); - const uint32_t C = viewmats.size(-3); - const uint32_t B = means.numel() / (N * 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(); @@ -64,25 +73,32 @@ projection_ewa_3dgs_fused_fwd( 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 &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; + 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) { + 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, + 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, @@ -96,12 +112,16 @@ projection_ewa_3dgs_fused_fwd( means2d.data_ptr(), depths.data_ptr(), conics.data_ptr(), - calc_compensations ? compensations.data_ptr() : nullptr + calc_compensations ? compensations.data_ptr() + : nullptr + ); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel ); - cgh.parallel_for(sycl::nd_range<1>(global_range, local_range), kernel); }); e.wait(); - }); + } + ); } return std::make_tuple(radii, means2d, depths, conics, compensations); diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp index 71012bd9..85418bf7 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -62,7 +62,9 @@ void launch_rasterize_2dgs_bwd_kernel( // Define the execution ranges sycl::range<3> localRange{1, tile_size, tile_size}; - sycl::range<3> globalRange{I, tile_height, tile_width}; + 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 @@ -155,7 +157,7 @@ rasterize_to_pixels_2dgs_bwd( 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] + const at::Tensor v_render_median, // [..., image_height, image_width] bool absgrad ) { // Check input tensors are contiguous @@ -197,7 +199,7 @@ rasterize_to_pixels_2dgs_bwd( // Launch kernel with appropriate dimension #define __GS__CALL_(DIM) \ case DIM: \ - launch_rasterize_2dgs_bwd_kernel( \ + 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, \ diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp index 3a7dc23b..37e21c3a 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp @@ -1,7 +1,7 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/RasterizeToPixels2DGSFwdKernel.hpp" namespace gsplat::xpu { @@ -11,20 +11,20 @@ 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, + 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, + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, // other params bool packed, uint32_t I, @@ -33,59 +33,77 @@ void launch_rasterize_2dgs_kernel( 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 + 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(); + 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_width}; + 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); - } - ); + // 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(); } @@ -107,7 +125,7 @@ rasterize_to_pixels_2dgs_fwd( 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] + const at::optional masks, // [..., tile_height, tile_width] // image size const uint32_t image_width, const uint32_t image_height, @@ -124,40 +142,86 @@ rasterize_to_pixels_2dgs_fwd( CHECK_CONTIGUOUS(normals); CHECK_CONTIGUOUS(tile_offsets); CHECK_CONTIGUOUS(flatten_ids); - if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); - if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); - + if (backgrounds.has_value()) + CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) + CHECK_CONTIGUOUS(masks.value()); + // Get dimensions bool packed = means2d.dim() == 2; - uint32_t N = packed ? 0 : means2d.size(-2); // number of gaussians - uint32_t I = tile_offsets.size(0); // number of images + 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 n_isects = flatten_ids.size(0); // number of intersections - uint32_t channels = colors.size(-1); // color dimension - + 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::Tensor renders = at::zeros({I, image_height, image_width, channels}, options_float); - at::Tensor alphas = at::zeros({I, image_height, image_width}, options_float); - at::Tensor render_normals = at::zeros({I, image_height, image_width, 3}, options_float); - at::Tensor render_distort = at::zeros({I, image_height, image_width}, options_float); - at::Tensor render_median = at::zeros({I, image_height, image_width}, options_float); - at::Tensor last_ids = at::zeros({I, image_height, image_width}, options_int); - at::Tensor median_ids = at::zeros({I, image_height, image_width}, options_int); + + 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 \ - ); \ +#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) { @@ -180,13 +244,20 @@ rasterize_to_pixels_2dgs_fwd( __GS__CALL_(257); __GS__CALL_(512); __GS__CALL_(513); - default: - TORCH_CHECK(false, "Unsupported number of channels: ", channels); + 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); + + return std::make_tuple( + renders, + alphas, + render_normals, + render_distort, + render_median, + last_ids, + median_ids + ); } } // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp index 1fd14693..35a6350d 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp @@ -103,12 +103,12 @@ void launch_rasterize_bwd_kernel( std::tuple rasterize_to_pixels_3dgs_bwd( // 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, + 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, @@ -117,11 +117,11 @@ rasterize_to_pixels_3dgs_bwd( const at::Tensor tile_offsets, const at::Tensor flatten_ids, // forward outputs - const at::Tensor render_alphas, - const at::Tensor last_ids, + 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, - const at::Tensor v_render_alphas, + 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 ) { @@ -138,11 +138,14 @@ rasterize_to_pixels_3dgs_bwd( if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + 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(1); + 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); diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp index a996ddfc..32a1b1a5 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp @@ -77,12 +77,12 @@ void launch_rasterize_kernel( std::tuple rasterize_to_pixels_3dgs_fwd( // 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, + 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, @@ -100,18 +100,32 @@ std::tuple rasterize_to_pixels_3dgs_fwd( if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + 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(1); + 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::Tensor renders = at::empty({C, image_height, image_width, channels}, options_float); - at::Tensor alphas = at::empty({C, image_height, image_width, 1}, options_float); - at::Tensor last_ids = at::empty({C, image_height, image_width}, options_int); + 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: \ diff --git a/gsplat/sycl/src/relocation.cpp b/gsplat/sycl/src/relocation.cpp index f51efa16..079d6655 100644 --- a/gsplat/sycl/src/relocation.cpp +++ b/gsplat/sycl/src/relocation.cpp @@ -1,11 +1,13 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" +#include "kernels/RelocationKernel.hpp" +#include "utils.hpp" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple relocation( at::Tensor opacities, // [N] at::Tensor scales, // [N, 3] @@ -13,7 +15,33 @@ std::tuple relocation( at::Tensor binoms, // [n_max, n_max] const int n_max ) { - throw std::runtime_error(std::string(__func__) + " is not implemented"); + 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 &q = + c10::xpu::getCurrentXPUStream().queue(); + q.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() + ) + ) + .wait(); + })); + + return std::make_tuple(new_opacities, new_scales); } -} // namespace gsplat::xpu \ No newline at end of file +} // namespace gsplat::xpu \ No newline at end of file diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 58752551..6c681f92 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -15,9 +15,6 @@ 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="Test requires CUDA backend" -) def expand(data: dict, batch_dims: Tuple[int, ...]): @@ -35,7 +32,7 @@ def expand(data: dict, batch_dims: Tuple[int, ...]): @pytest.fixture -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@requires_backend def test_data(): C = 3 N = 1000 @@ -138,7 +135,7 @@ 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") +@pytest.mark.skipif(device is None, reason="No GPU device") @pytest.mark.parametrize("sparse_grad", [False]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_fully_fused_projection_packed_2dgs( @@ -262,14 +259,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._torch_impl_2dgs import _rasterize_to_pixels_2dgs - from gsplat.cuda._wrapper import ( + from gsplat import ( fully_fused_projection_2dgs, isect_offset_encode, isect_tiles, @@ -341,19 +338,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) @@ -373,31 +371,32 @@ 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 52a01528..77abfe25 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -26,7 +26,7 @@ device = None requires_backend = pytest.mark.skipif( - gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" + 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" @@ -555,20 +555,22 @@ 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) + + 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_render_colors, v_render_alphas = torch.randn_like( render_colors @@ -578,13 +580,15 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. + (render_alphas * v_render_alphas).sum(), (means2d, conics, colors, opacities, backgrounds), ) - _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) + + 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..ed410d3d 100644 --- a/tests/test_rasterization.py +++ b/tests/test_rasterization.py @@ -11,10 +11,20 @@ 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 +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL XPU backend available" +) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +@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 +91,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..24599f21 100644 --- a/tests/test_strategy.py +++ b/tests/test_strategy.py @@ -8,11 +8,21 @@ import pytest import torch - -device = torch.device("cuda:0") - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="No CUDA device") +import gsplat +if gsplat.BACKEND == "sycl": + device = torch.device("xpu:0") + torch_acc = torch.xpu +elif gsplat.BACKEND == "cuda": + device = torch.device("cuda:0") + torch_acc = torch.cuda +else: + device = None + +requires_backend = pytest.mark.skipif( + gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" +) + +@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 From 6524660b286c992ce58618066c13d8ee47d523e3 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Sun, 16 Nov 2025 00:09:43 -0800 Subject: [PATCH 29/56] CI tests and docs should pass now. black formatting ALways import functions in __init__ even if no backend, so that docs are built in CI and tests run in CI (all skipped). Skip sycl compiler check in setup.py - cmake has more robust path search and it will complain if not found. --- .github/workflows/core_tests.yml | 2 +- .github/workflows/doc.yml | 2 +- docs/source/apis/utils.rst | 4 - examples/gsplat_viewer.py | 24 +++--- examples/gsplat_viewer_2dgs.py | 24 +++--- gsplat/__init__.py | 117 +++++++++++--------------- gsplat/compression/png_compression.py | 4 +- gsplat/cuda/_wrapper.py | 20 +++-- gsplat/distributed.py | 27 +++--- gsplat/rendering.py | 14 ++- gsplat/strategy/default.py | 1 + gsplat/strategy/ops.py | 11 ++- gsplat/sycl/_wrapper.py | 26 ++++-- profiling/main.py | 2 +- setup.py | 11 +-- tests/test_2dgs.py | 24 ++++-- tests/test_basic.py | 73 +++++++++------- tests/test_rasterization.py | 7 +- tests/test_strategy.py | 4 +- 19 files changed, 197 insertions(+), 200 deletions(-) diff --git a/.github/workflows/core_tests.yml b/.github/workflows/core_tests.yml index a4db904f..ae6330cd 100644 --- a/.github/workflows/core_tests.yml +++ b/.github/workflows/core_tests.yml @@ -26,7 +26,7 @@ jobs: 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 . + 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..cf70066f 100644 --- a/.github/workflows/doc.yml +++ b/.github/workflows/doc.yml @@ -27,7 +27,7 @@ jobs: 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 . + BUILD_NO_CUDA=1 pip install --no-build-isolation . # Get version. - name: Get version + subdirectory 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/gsplat_viewer.py b/examples/gsplat_viewer.py index e47d75a8..f2290e59 100644 --- a/examples/gsplat_viewer.py +++ b/examples/gsplat_viewer.py @@ -17,14 +17,14 @@ class GsplatRenderTabState(RenderTabState): radius_clip: float = 0.0 eps2d: float = 0.3 backgrounds: Tuple[float, float, float] = (0.0, 0.0, 0.0) - render_mode: Literal[ - "rgb", "depth(accumulated)", "depth(expected)", "alpha" - ] = "rgb" + render_mode: Literal["rgb", "depth(accumulated)", "depth(expected)", "alpha"] = ( + "rgb" + ) normalize_nearfar: bool = False inverse: bool = False - colormap: Literal[ - "turbo", "viridis", "magma", "inferno", "cividis", "gray" - ] = "turbo" + colormap: Literal["turbo", "viridis", "magma", "inferno", "cividis", "gray"] = ( + "turbo" + ) rasterize_mode: Literal["classic", "antialiased"] = "classic" camera_model: Literal["pinhole", "ortho", "fisheye"] = "pinhole" @@ -239,9 +239,9 @@ def _(_) -> None: def _after_render(self): # Update the GUI elements with current values - self._rendering_tab_handles[ - "total_gs_count_number" - ].value = self.render_tab_state.total_gs_count - self._rendering_tab_handles[ - "rendered_gs_count_number" - ].value = self.render_tab_state.rendered_gs_count + self._rendering_tab_handles["total_gs_count_number"].value = ( + self.render_tab_state.total_gs_count + ) + self._rendering_tab_handles["rendered_gs_count_number"].value = ( + self.render_tab_state.rendered_gs_count + ) diff --git a/examples/gsplat_viewer_2dgs.py b/examples/gsplat_viewer_2dgs.py index d9b258cd..675fb339 100644 --- a/examples/gsplat_viewer_2dgs.py +++ b/examples/gsplat_viewer_2dgs.py @@ -17,14 +17,14 @@ class GsplatRenderTabState(RenderTabState): radius_clip: float = 0.0 eps2d: float = 0.3 backgrounds: Tuple[float, float, float] = (0.0, 0.0, 0.0) - render_mode: Literal[ - "rgb", "depth(accumulated)", "depth(expected)", "alpha" - ] = "rgb" + render_mode: Literal["rgb", "depth(accumulated)", "depth(expected)", "alpha"] = ( + "rgb" + ) normalize_nearfar: bool = False inverse: bool = False - colormap: Literal[ - "turbo", "viridis", "magma", "inferno", "cividis", "gray" - ] = "turbo" + colormap: Literal["turbo", "viridis", "magma", "inferno", "cividis", "gray"] = ( + "turbo" + ) class GsplatViewer(Viewer): @@ -211,9 +211,9 @@ def _(_) -> None: def _after_render(self): # Update the GUI elements with current values - self._rendering_tab_handles[ - "total_gs_count_number" - ].value = self.render_tab_state.total_gs_count - self._rendering_tab_handles[ - "rendered_gs_count_number" - ].value = self.render_tab_state.rendered_gs_count + self._rendering_tab_handles["total_gs_count_number"].value = ( + self.render_tab_state.total_gs_count + ) + self._rendering_tab_handles["rendered_gs_count_number"].value = ( + self.render_tab_state.rendered_gs_count + ) diff --git a/gsplat/__init__.py b/gsplat/__init__.py index 527c77ad..00bafdfa 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -1,80 +1,63 @@ import os import sys import torch -import warnings BACKEND: str = "" +torch_acc = torch.cpu +_force_backend = os.getenv("GSPLAT_BACKEND", "").lower() -FORCE_BACKEND = os.getenv("GSPLAT_BACKEND", "").lower() - -if FORCE_BACKEND == "cuda" or (FORCE_BACKEND == "" and torch.cuda.is_available()): - try: - BACKEND = "cuda" - from .cuda._wrapper import ( - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - proj, - quat_scale_to_covar_preci, - rasterize_to_indices_in_range, - rasterize_to_indices_in_range_2dgs, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - world_to_cam, - ) - - torch_acc = torch.cuda - print("gsplat: CUDA backend successfully loaded.", file=sys.stderr) - except ImportError: - if FORCE_BACKEND == "cuda": - print( - "gsplat: Error! GSPLAT_BACKEND=cuda was set but CUDA backend failed to load.", - file=sys.stderr, - ) - pass - -if not BACKEND and (FORCE_BACKEND in ("sycl", "xpu") or FORCE_BACKEND == "" and torch.xpu.is_available()): - try: - BACKEND = "sycl" - from .sycl._wrapper import ( - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - proj, - quat_scale_to_covar_preci, - rasterize_to_indices_in_range, - rasterize_to_indices_in_range_2dgs, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - world_to_cam, - ) +from .cuda._wrapper import ( # Default to CUDA imports, works even if no CUDA is available + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + proj, + quat_scale_to_covar_preci, + rasterize_to_indices_in_range, + rasterize_to_indices_in_range_2dgs, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + world_to_cam, +) - torch_acc = torch.xpu - print("gsplat: SYCL XPU backend successfully loaded.", file=sys.stderr) - except ImportError as e: - if FORCE_BACKEND in ("sycl", "xpu"): - print( - f"gsplat: Error! GSPLAT_BACKEND={FORCE_BACKEND} was set but SYCL XPU backend failed to load: {e}", - file=sys.stderr, - ) - pass +if _force_backend == "cuda" or (_force_backend == "" and torch.cuda.is_available()): + BACKEND = "cuda" + torch_acc = torch.cuda + print("gsplat: Using CUDA backend.", file=sys.stderr) + # Functions already imported above -if not BACKEND: - print( - "gsplat: Warning! No high-performance backend (CUDA or SYCL XPU) found.", - file=sys.stderr, +if ( + not BACKEND + and _force_backend in ("sycl", "xpu") + or _force_backend == "" + and torch.xpu.is_available() +): + from .sycl._wrapper import ( # Overwrite imports for SYCL backend + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + proj, + quat_scale_to_covar_preci, + rasterize_to_indices_in_range, + rasterize_to_indices_in_range_2dgs, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + world_to_cam, ) + BACKEND = "sycl" + torch_acc = torch.xpu + print("gsplat: Using SYCL XPU backend.", file=sys.stderr) + from .compression import PngCompression from .exporter import export_splats diff --git a/gsplat/compression/png_compression.py b/gsplat/compression/png_compression.py index 0c5010cc..046a666a 100644 --- a/gsplat/compression/png_compression.py +++ b/gsplat/compression/png_compression.py @@ -368,9 +368,7 @@ def _compress_kmeans( maxs = torch.max(centroids) centroids_norm = (centroids - mins) / (maxs - mins) centroids_norm = centroids_norm.detach().cpu().numpy() - centroids_quant = ( - (centroids_norm * (2**quantization - 1)).round().astype(np.uint8) - ) + centroids_quant = (centroids_norm * (2**quantization - 1)).round().astype(np.uint8) labels = labels.astype(np.uint16) npz_dict = { diff --git a/gsplat/cuda/_wrapper.py b/gsplat/cuda/_wrapper.py index 74b01ca5..69890c36 100644 --- a/gsplat/cuda/_wrapper.py +++ b/gsplat/cuda/_wrapper.py @@ -317,7 +317,7 @@ 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. @@ -1239,9 +1239,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 @@ -1509,9 +1511,13 @@ 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( - "rasterize_to_pixels_from_world_3dgs_bwd" - )( + ( + v_means, + v_quats, + v_scales, + v_colors, + v_opacities, + ) = _make_lazy_cuda_func("rasterize_to_pixels_from_world_3dgs_bwd")( means, quats, scales, diff --git a/gsplat/distributed.py b/gsplat/distributed.py index 26e2de23..9b02bb52 100644 --- a/gsplat/distributed.py +++ b/gsplat/distributed.py @@ -6,14 +6,16 @@ import torch.distributed.nn.functional as distF from torch import Tensor -import gsplat -from gsplat import torch_acc +from . import torch_acc, BACKEND def _get_distributed_backend(): - if gsplat.BACKEND == "sycl": - return "ccl" - return "nccl" + if BACKEND == "sycl": + return "xccl" + elif BACKEND == "cuda": + return "nccl" + else: + return "gloo" def all_gather_int32( @@ -45,10 +47,10 @@ def all_gather_int32( value_tensor = torch.tensor(value, dtype=torch.int, device=device) else: value_tensor = value - - if gsplat.BACKEND == "cuda": + + if BACKEND == "cuda": assert value_tensor.is_cuda, "value should be on CUDA" - elif gsplat.BACKEND == "sycl": + elif BACKEND == "sycl": assert value_tensor.is_xpu, "value should be on XPU" # gather @@ -333,12 +335,6 @@ def fn(local_rank: int, world_rank: int, world_size: int, args: Any) -> None: cli(fn, None, verbose=True) ``` """ - if gsplat.BACKEND == "cuda": - assert torch.cuda.is_available(), "CUDA device is required!" - elif gsplat.BACKEND == "sycl": - import torch.xpu - assert torch.xpu.is_available(), "XPU 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() @@ -347,7 +343,6 @@ 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_acc.device_count() distributed = world_size > 1 @@ -377,4 +372,4 @@ def fn(local_rank: int, world_rank: int, world_size: int, args: Any) -> None: print("process " + str(i) + " finished") return True else: - return _distributed_worker(0, 1, fn=fn, args=args) \ No newline at end of file + return _distributed_worker(0, 1, fn=fn, args=args) diff --git a/gsplat/rendering.py b/gsplat/rendering.py index 3e88be4e..8a839581 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -10,8 +10,8 @@ from . import BACKEND # Now, conditionally import the functions based on the detected backend. -if BACKEND == "cuda": - from .cuda._wrapper import ( +if BACKEND == "sycl": + from .sycl._wrapper import ( RollingShutterType, fully_fused_projection, fully_fused_projection_2dgs, @@ -23,8 +23,8 @@ rasterize_to_pixels_eval3d, spherical_harmonics, ) -elif BACKEND == "sycl": - from .sycl._wrapper import ( +else: # CUDA or no backend (e.g., CPU only for docs and testing) + from .cuda._wrapper import ( RollingShutterType, fully_fused_projection, fully_fused_projection_2dgs, @@ -36,10 +36,6 @@ rasterize_to_pixels_eval3d, spherical_harmonics, ) -else: - # If no backend is found, you can either raise an error or define dummy functions - # to avoid crashing, depending on your needs. - raise ImportError("gsplat: No backend (CUDA or SYCL) found, cannot import backend-specific functions.") from .distributed import ( all_gather_int32, @@ -83,7 +79,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 = None, + ftheta_coeffs=None, # rolling shutter rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] diff --git a/gsplat/strategy/default.py b/gsplat/strategy/default.py index 61f68aee..8b3a6b18 100644 --- a/gsplat/strategy/default.py +++ b/gsplat/strategy/default.py @@ -8,6 +8,7 @@ from .ops import duplicate, remove, reset_opa, split from .. import torch_acc + @dataclass class DefaultStrategy(Strategy): """A default strategy that follows the original 3DGS paper: diff --git a/gsplat/strategy/ops.py b/gsplat/strategy/ops.py index d14b37e2..7a7aa66d 100644 --- a/gsplat/strategy/ops.py +++ b/gsplat/strategy/ops.py @@ -6,16 +6,15 @@ from torch import Tensor from gsplat import BACKEND -if BACKEND == "cuda": - from gsplat.cuda._wrapper import ( + +if BACKEND == "sycl": + from gsplat.sycl._wrapper import ( quat_scale_to_covar_preci, ) -elif BACKEND == "sycl": - from gsplat.sycl._wrapper import ( +else: # BACKEND == "cuda" or None + from gsplat.cuda._wrapper import ( quat_scale_to_covar_preci, ) -else: - raise ImportError("gsplat: No backend loaded, cannot import strategy ops.") from gsplat.relocation import compute_relocation from gsplat.utils import normalized_quat_to_rotmat diff --git a/gsplat/sycl/_wrapper.py b/gsplat/sycl/_wrapper.py index 0f32cb01..8c247cef 100644 --- a/gsplat/sycl/_wrapper.py +++ b/gsplat/sycl/_wrapper.py @@ -11,17 +11,21 @@ def _make_lazy_sycl_func(name: str) -> Callable: """Creates a lazy-loading function for the SYCL backend.""" + def call_sycl(*args, **kwargs): # pylint: disable=import-outside-toplevel from ._backend import _C + return getattr(_C, name)(*args, **kwargs) + return call_sycl + def _make_lazy_sycl_obj(name: str) -> Any: """Creates a lazy-loading object accessor for the SYCL backend.""" # pylint: disable=import-outside-toplevel from ._backend import _C - + obj = _C for name_split in name.split("."): obj = getattr(obj, name_split) @@ -316,7 +320,7 @@ 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. @@ -1238,9 +1242,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 @@ -1508,9 +1514,13 @@ def backward( tile_size = ctx.tile_size ftheta_coeffs = ctx.ftheta_coeffs - (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_sycl_func( - "rasterize_to_pixels_from_world_3dgs_bwd" - )( + ( + v_means, + v_quats, + v_scales, + v_colors, + v_opacities, + ) = _make_lazy_sycl_func("rasterize_to_pixels_from_world_3dgs_bwd")( means, quats, scales, diff --git a/profiling/main.py b/profiling/main.py index 1e156d84..028295f0 100644 --- a/profiling/main.py +++ b/profiling/main.py @@ -11,7 +11,7 @@ import torch from typing_extensions import Callable, Literal -from gsplat import torch_acc, BACKEND +from gsplat import torch_acc, BACKEND from gsplat._helper import load_test_data from gsplat.distributed import cli from gsplat.rendering import rasterization diff --git a/setup.py b/setup.py index c258ea91..223afbcb 100644 --- a/setup.py +++ b/setup.py @@ -29,16 +29,7 @@ except (ImportError, AttributeError): pass -has_sycl_compiler = False -if ( - sp.run(["icpx", "--version"], stdout=sp.DEVNULL, stderr=sp.DEVNULL).returncode == 0 -) or ( - sp.run(["dpcpp", "--version"], stdout=sp.DEVNULL, stderr=sp.DEVNULL).returncode == 0 -): - has_sycl_compiler = True - -BUILD_SYCL = has_xpu and has_sycl_compiler - +BUILD_SYCL = has_xpu BUILD_NO_CUDA = os.getenv("BUILD_NO_CUDA", "0") == "1" WITH_SYMBOLS = os.getenv("WITH_SYMBOLS", "0") == "1" LINE_INFO = os.getenv("LINE_INFO", "0") == "1" diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 6c681f92..241e36e1 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -4,13 +4,14 @@ import torch from typing_extensions import Tuple -import gsplat +import gsplat + if gsplat.BACKEND == "sycl": device = torch.device("xpu:0") elif gsplat.BACKEND == "cuda": device = torch.device("cuda:0") else: - device = None + device = torch.device("cpu") requires_backend = pytest.mark.skipif( gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" @@ -32,7 +33,6 @@ def expand(data: dict, batch_dims: Tuple[int, ...]): @pytest.fixture -@requires_backend def test_data(): C = 3 N = 1000 @@ -135,7 +135,7 @@ 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(device is None, reason="No GPU device") +@requires_backend @pytest.mark.parametrize("sparse_grad", [False]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_fully_fused_projection_packed_2dgs( @@ -322,7 +322,13 @@ def test_rasterize_to_pixels_2dgs( normals.requires_grad = True densify.requires_grad = True - (render_colors, render_alphas, render_normals, _, _,) = rasterize_to_pixels_2dgs( + ( + render_colors, + render_alphas, + render_normals, + _, + _, + ) = rasterize_to_pixels_2dgs( means2d, ray_transforms, colors, @@ -338,7 +344,7 @@ def test_rasterize_to_pixels_2dgs( distloss=True, ) - if gsplat.BACKEND != "sycl": # nerfacc required for comparison + if gsplat.BACKEND != "sycl": # nerfacc required for comparison _render_colors, _render_alphas, _render_normals = _rasterize_to_pixels_2dgs( means2d, ray_transforms, @@ -371,7 +377,7 @@ def test_rasterize_to_pixels_2dgs( (means2d, ray_transforms, colors, opacities, backgrounds, normals), ) - if gsplat.BACKEND != "sycl": # nerfacc required for comparison + if gsplat.BACKEND != "sycl": # nerfacc required for comparison ( _v_means2d, _v_ray_transforms, @@ -389,7 +395,9 @@ def test_rasterize_to_pixels_2dgs( # 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) + 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) diff --git a/tests/test_basic.py b/tests/test_basic.py index 77abfe25..ef9f3c26 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -26,7 +26,8 @@ device = None requires_backend = pytest.mark.skipif( - gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL XPU backend available" + 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" @@ -313,20 +314,24 @@ def test_fully_fused_projection_packed( calc_compensations=calc_compensations, camera_model=camera_model, ) - _radii, _means2d, _depths, _conics, _compensations = ( - gsplat.fully_fused_projection( - means, - None, - quats, - scales, - viewmats, - Ks, - width, - height, - packed=False, - calc_compensations=calc_compensations, - camera_model=camera_model, - ) + ( + _radii, + _means2d, + _depths, + _conics, + _compensations, + ) = gsplat.fully_fused_projection( + means, + None, + quats, + scales, + viewmats, + Ks, + width, + height, + packed=False, + calc_compensations=calc_compensations, + camera_model=camera_model, ) else: covars, _ = gsplat.quat_scale_to_covar_preci(quats, scales, triu=True) @@ -344,20 +349,24 @@ def test_fully_fused_projection_packed( calc_compensations=calc_compensations, camera_model=camera_model, ) - _radii, _means2d, _depths, _conics, _compensations = ( - gsplat.fully_fused_projection( - means, - covars, - None, - None, - viewmats, - Ks, - width, - height, - packed=False, - calc_compensations=calc_compensations, - camera_model=camera_model, - ) + ( + _radii, + _means2d, + _depths, + _conics, + _compensations, + ) = gsplat.fully_fused_projection( + means, + covars, + None, + None, + viewmats, + Ks, + width, + height, + packed=False, + calc_compensations=calc_compensations, + camera_model=camera_model, ) ( @@ -556,7 +565,7 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. backgrounds=backgrounds, ) - if gsplat.BACKEND != "sycl": # nerfacc required for comparison + if gsplat.BACKEND != "sycl": # nerfacc required for comparison _render_colors, _render_alphas = _rasterize_to_pixels( means2d, conics, @@ -580,8 +589,8 @@ def test_rasterize_to_pixels(test_data, channels: int, batch_dims: Tuple[int, .. + (render_alphas * v_render_alphas).sum(), (means2d, conics, colors, opacities, backgrounds), ) - - if gsplat.BACKEND != "sycl": # nerfacc required for comparison + + if gsplat.BACKEND != "sycl": # nerfacc required for comparison _grads = torch.autograd.grad( (_render_colors * v_render_colors).sum() + (_render_alphas * v_render_alphas).sum(), diff --git a/tests/test_rasterization.py b/tests/test_rasterization.py index ed410d3d..df53c5af 100644 --- a/tests/test_rasterization.py +++ b/tests/test_rasterization.py @@ -13,6 +13,7 @@ # device = torch.device("cuda:0") import gsplat + if gsplat.BACKEND == "sycl": device = torch.device("xpu:0") elif gsplat.BACKEND == "cuda": @@ -21,9 +22,11 @@ device = None requires_backend = pytest.mark.skipif( - gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL XPU backend available" + 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]) @@ -91,7 +94,7 @@ def test_rasterization( elif render_mode == "RGB+D": assert renders.shape == batch_dims + (C, height, width, 4) - if gsplat.BACKEND != "sycl": # nerfacc required for comparison + if gsplat.BACKEND != "sycl": # nerfacc required for comparison _renders, _alphas, _meta = _rasterization( means=means, quats=quats, diff --git a/tests/test_strategy.py b/tests/test_strategy.py index 24599f21..9e48cd1b 100644 --- a/tests/test_strategy.py +++ b/tests/test_strategy.py @@ -8,7 +8,8 @@ import pytest import torch -import gsplat +import gsplat + if gsplat.BACKEND == "sycl": device = torch.device("xpu:0") torch_acc = torch.xpu @@ -22,6 +23,7 @@ gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" ) + @requires_backend def test_strategy(): from gsplat.rendering import rasterization From e261c0480a85a29fa8ec8f2f6a91fd2891af2edb Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Sun, 16 Nov 2025 06:09:36 -0800 Subject: [PATCH 30/56] Use pytorch 2.6 for CI (docs and core_tests) check is torch has xpu --- .github/workflows/core_tests.yml | 6 +++--- .github/workflows/doc.yml | 2 +- gsplat/__init__.py | 1 + setup.py | 2 +- tests/test_strategy.py | 4 +--- 5 files changed, 7 insertions(+), 8 deletions(-) diff --git a/.github/workflows/core_tests.yml b/.github/workflows/core_tests.yml index ae6330cd..97fde27f 100644 --- a/.github/workflows/core_tests.yml +++ b/.github/workflows/core_tests.yml @@ -18,14 +18,14 @@ 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 + 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 diff --git a/.github/workflows/doc.yml b/.github/workflows/doc.yml index cf70066f..31e97142 100644 --- a/.github/workflows/doc.yml +++ b/.github/workflows/doc.yml @@ -26,7 +26,7 @@ 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 + 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. diff --git a/gsplat/__init__.py b/gsplat/__init__.py index 00bafdfa..9df5004f 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -34,6 +34,7 @@ not BACKEND and _force_backend in ("sycl", "xpu") or _force_backend == "" + and hasattr(torch, "xpu") and torch.xpu.is_available() ): from .sycl._wrapper import ( # Overwrite imports for SYCL backend diff --git a/setup.py b/setup.py index 223afbcb..987bc1ee 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,7 @@ try: import torch - has_xpu = torch.xpu.is_available() + has_xpu = has_cuda and hasattr(torch, 'xpu') and torch.xpu.is_available() except (ImportError, AttributeError): pass diff --git a/tests/test_strategy.py b/tests/test_strategy.py index 9e48cd1b..a5800d7d 100644 --- a/tests/test_strategy.py +++ b/tests/test_strategy.py @@ -12,12 +12,10 @@ if gsplat.BACKEND == "sycl": device = torch.device("xpu:0") - torch_acc = torch.xpu elif gsplat.BACKEND == "cuda": device = torch.device("cuda:0") - torch_acc = torch.cuda else: - device = None + device = torch.device("cpu") requires_backend = pytest.mark.skipif( gsplat.BACKEND not in ("cuda", "sycl"), reason="No CUDA or SYCL backend available" From d03c5f3442e4e87523d70ede3be4e120bff16b52 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Sun, 16 Nov 2025 06:54:05 -0800 Subject: [PATCH 31/56] black v22 and clang-format --- examples/gsplat_viewer.py | 24 +- examples/gsplat_viewer_2dgs.py | 24 +- gsplat/compression/png_compression.py | 4 +- gsplat/cuda/_wrapper.py | 10 +- gsplat/sycl/_wrapper.py | 10 +- gsplat/sycl/include/Cameras.h | 6 +- gsplat/sycl/include/Common.h | 3 +- gsplat/sycl/include/Ops.h | 71 +- gsplat/sycl/include/gsplat_sycl_utils.hpp | 72 +- gsplat/sycl/include/helpers.hpp | 14 +- .../include/kernels/ComputeShBwdKernel.hpp | 46 +- .../include/kernels/ComputeShFwdKernel.hpp | 32 +- .../kernels/FullyFusedProjectionBwdKernel.hpp | 243 +++--- .../kernels/FullyFusedProjectionFwdKernel.hpp | 181 ++-- .../kernels/IsectOffsetEncodeKernel.hpp | 28 +- .../sycl/include/kernels/IsectTilesKernel.hpp | 107 +-- .../kernels/PackedProjectionBwdKernel.hpp | 333 +++++--- .../kernels/PackedProjectionFwdKernel.hpp | 253 ++++-- gsplat/sycl/include/kernels/ProjBwdKernel.hpp | 172 ++-- gsplat/sycl/include/kernels/ProjFwdKernel.hpp | 85 +- .../kernels/Projection2DGSFusedBwdKernel.hpp | 209 +++-- .../kernels/Projection2DGSFusedFwdKernel.hpp | 160 ++-- .../QuatScaleToCovarPreciBwdKernel.hpp | 62 +- .../QuatScaleToCovarPreciFwdKernel.hpp | 60 +- .../RasterizeToPixels2DGSBwdKernel.hpp | 607 ++++++++------ .../RasterizeToPixels2DGSFwdKernel.hpp | 325 ++++---- .../kernels/RasterizeToPixelsBwdKernel.hpp | 780 ++++++++++-------- .../kernels/RasterizeToPixelsFwdKernel.hpp | 564 +++++++------ .../sycl/include/kernels/RelocationKernel.hpp | 4 +- .../include/kernels/WorldToCamBwdKernel.hpp | 90 +- .../include/kernels/WorldToCamFwdKernel.hpp | 59 +- gsplat/sycl/include/proj.hpp | 4 +- gsplat/sycl/include/quat.hpp | 5 +- .../include/quat_scale_to_covar_preci.hpp | 3 +- gsplat/sycl/include/spherical_harmonics.hpp | 2 - gsplat/sycl/include/types.hpp | 2 - gsplat/sycl/include/utils.hpp | 19 +- gsplat/sycl/src/adam.cpp | 8 +- gsplat/sycl/src/intersect_offset.cpp | 38 +- gsplat/sycl/src/intersect_tile.cpp | 66 +- gsplat/sycl/src/null.cpp | 8 +- gsplat/sycl/src/projection_2dgs_fused_bwd.cpp | 47 +- gsplat/sycl/src/projection_2dgs_fused_fwd.cpp | 47 +- .../sycl/src/projection_2dgs_packed_bwd.cpp | 14 +- .../sycl/src/projection_2dgs_packed_fwd.cpp | 8 +- .../src/projection_ewa_3dgs_fused_bwd.cpp | 79 +- .../src/projection_ewa_3dgs_fused_fwd.cpp | 4 +- .../src/projection_ewa_3dgs_packed_bwd.cpp | 112 ++- .../src/projection_ewa_3dgs_packed_fwd.cpp | 218 +++-- gsplat/sycl/src/projection_ewa_simple_bwd.cpp | 63 +- gsplat/sycl/src/projection_ewa_simple_fwd.cpp | 61 +- gsplat/sycl/src/projection_ut_3dgs_fused.cpp | 20 +- .../src/quat_scale_to_covar_preci_bwd.cpp | 46 +- .../src/quat_scale_to_covar_preci_fwd.cpp | 42 +- gsplat/sycl/src/rasterize_to_indices_2dgs.cpp | 8 +- gsplat/sycl/src/rasterize_to_indices_3dgs.cpp | 8 +- .../sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 283 ++++--- .../sycl/src/rasterize_to_pixels_3dgs_bwd.cpp | 233 ++++-- .../sycl/src/rasterize_to_pixels_3dgs_fwd.cpp | 169 ++-- ...asterize_to_pixels_from_world_3dgs_bwd.cpp | 30 +- ...asterize_to_pixels_from_world_3dgs_fwd.cpp | 22 +- gsplat/sycl/src/relocation.cpp | 30 +- gsplat/sycl/src/spherical_harmonics_bwd.cpp | 46 +- gsplat/sycl/src/spherical_harmonics_fwd.cpp | 58 +- setup.py | 2 +- tests/test_2dgs.py | 8 +- 66 files changed, 3668 insertions(+), 2783 deletions(-) diff --git a/examples/gsplat_viewer.py b/examples/gsplat_viewer.py index f2290e59..e47d75a8 100644 --- a/examples/gsplat_viewer.py +++ b/examples/gsplat_viewer.py @@ -17,14 +17,14 @@ class GsplatRenderTabState(RenderTabState): radius_clip: float = 0.0 eps2d: float = 0.3 backgrounds: Tuple[float, float, float] = (0.0, 0.0, 0.0) - render_mode: Literal["rgb", "depth(accumulated)", "depth(expected)", "alpha"] = ( - "rgb" - ) + render_mode: Literal[ + "rgb", "depth(accumulated)", "depth(expected)", "alpha" + ] = "rgb" normalize_nearfar: bool = False inverse: bool = False - colormap: Literal["turbo", "viridis", "magma", "inferno", "cividis", "gray"] = ( - "turbo" - ) + colormap: Literal[ + "turbo", "viridis", "magma", "inferno", "cividis", "gray" + ] = "turbo" rasterize_mode: Literal["classic", "antialiased"] = "classic" camera_model: Literal["pinhole", "ortho", "fisheye"] = "pinhole" @@ -239,9 +239,9 @@ def _(_) -> None: def _after_render(self): # Update the GUI elements with current values - self._rendering_tab_handles["total_gs_count_number"].value = ( - self.render_tab_state.total_gs_count - ) - self._rendering_tab_handles["rendered_gs_count_number"].value = ( - self.render_tab_state.rendered_gs_count - ) + self._rendering_tab_handles[ + "total_gs_count_number" + ].value = self.render_tab_state.total_gs_count + self._rendering_tab_handles[ + "rendered_gs_count_number" + ].value = self.render_tab_state.rendered_gs_count diff --git a/examples/gsplat_viewer_2dgs.py b/examples/gsplat_viewer_2dgs.py index 675fb339..d9b258cd 100644 --- a/examples/gsplat_viewer_2dgs.py +++ b/examples/gsplat_viewer_2dgs.py @@ -17,14 +17,14 @@ class GsplatRenderTabState(RenderTabState): radius_clip: float = 0.0 eps2d: float = 0.3 backgrounds: Tuple[float, float, float] = (0.0, 0.0, 0.0) - render_mode: Literal["rgb", "depth(accumulated)", "depth(expected)", "alpha"] = ( - "rgb" - ) + render_mode: Literal[ + "rgb", "depth(accumulated)", "depth(expected)", "alpha" + ] = "rgb" normalize_nearfar: bool = False inverse: bool = False - colormap: Literal["turbo", "viridis", "magma", "inferno", "cividis", "gray"] = ( - "turbo" - ) + colormap: Literal[ + "turbo", "viridis", "magma", "inferno", "cividis", "gray" + ] = "turbo" class GsplatViewer(Viewer): @@ -211,9 +211,9 @@ def _(_) -> None: def _after_render(self): # Update the GUI elements with current values - self._rendering_tab_handles["total_gs_count_number"].value = ( - self.render_tab_state.total_gs_count - ) - self._rendering_tab_handles["rendered_gs_count_number"].value = ( - self.render_tab_state.rendered_gs_count - ) + self._rendering_tab_handles[ + "total_gs_count_number" + ].value = self.render_tab_state.total_gs_count + self._rendering_tab_handles[ + "rendered_gs_count_number" + ].value = self.render_tab_state.rendered_gs_count diff --git a/gsplat/compression/png_compression.py b/gsplat/compression/png_compression.py index 046a666a..0c5010cc 100644 --- a/gsplat/compression/png_compression.py +++ b/gsplat/compression/png_compression.py @@ -368,7 +368,9 @@ def _compress_kmeans( maxs = torch.max(centroids) centroids_norm = (centroids - mins) / (maxs - mins) centroids_norm = centroids_norm.detach().cpu().numpy() - centroids_quant = (centroids_norm * (2**quantization - 1)).round().astype(np.uint8) + centroids_quant = ( + (centroids_norm * (2**quantization - 1)).round().astype(np.uint8) + ) labels = labels.astype(np.uint16) npz_dict = { diff --git a/gsplat/cuda/_wrapper.py b/gsplat/cuda/_wrapper.py index 69890c36..50ba0f03 100644 --- a/gsplat/cuda/_wrapper.py +++ b/gsplat/cuda/_wrapper.py @@ -1511,13 +1511,9 @@ 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("rasterize_to_pixels_from_world_3dgs_bwd")( + (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_cuda_func( + "rasterize_to_pixels_from_world_3dgs_bwd" + )( means, quats, scales, diff --git a/gsplat/sycl/_wrapper.py b/gsplat/sycl/_wrapper.py index 8c247cef..8e188e8b 100644 --- a/gsplat/sycl/_wrapper.py +++ b/gsplat/sycl/_wrapper.py @@ -1514,13 +1514,9 @@ def backward( tile_size = ctx.tile_size ftheta_coeffs = ctx.ftheta_coeffs - ( - v_means, - v_quats, - v_scales, - v_colors, - v_opacities, - ) = _make_lazy_sycl_func("rasterize_to_pixels_from_world_3dgs_bwd")( + (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_sycl_func( + "rasterize_to_pixels_from_world_3dgs_bwd" + )( means, quats, scales, diff --git a/gsplat/sycl/include/Cameras.h b/gsplat/sycl/include/Cameras.h index f99a937c..bd3d38ed 100644 --- a/gsplat/sycl/include/Cameras.h +++ b/gsplat/sycl/include/Cameras.h @@ -51,8 +51,10 @@ struct FThetaCameraDistortionParameters { ANGLE_TO_PIXELDIST, }; PolynomialType reference_poly; - std::array pixeldist_to_angle_poly; // backward polynomial - std::array angle_to_pixeldist_poly; // forward polynomial + 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 index 407ae2ca..d92a978d 100644 --- a/gsplat/sycl/include/Common.h +++ b/gsplat/sycl/include/Common.h @@ -4,7 +4,7 @@ #include #include -namespace gsplat::xpu { +namespace gsplat::xpu { // // Some Macros. @@ -16,7 +16,6 @@ namespace gsplat::xpu { CHECK_XPU(x); \ CHECK_CONTIGUOUS(x) - // // Legacy Camera Types // diff --git a/gsplat/sycl/include/Ops.h b/gsplat/sycl/include/Ops.h index bb8bb37c..80ec7e8a 100644 --- a/gsplat/sycl/include/Ops.h +++ b/gsplat/sycl/include/Ops.h @@ -1,13 +1,13 @@ // A collection of operators for gsplat #pragma once -#include -#include #include "Cameras.h" #include "Common.h" #include "types.hpp" +#include +#include -namespace gsplat::xpu { +namespace gsplat::xpu { // null operator for tutorial. Does nothing. at::Tensor null(const at::Tensor input); @@ -228,7 +228,7 @@ std::tuple rasterize_to_pixels_3dgs_fwd( 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] + const at::optional masks, // [..., tile_height, tile_width] // image size const uint32_t image_width, const uint32_t image_height, @@ -245,7 +245,7 @@ rasterize_to_pixels_3dgs_bwd( 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] + const at::optional masks, // [..., tile_height, tile_width] // image size const uint32_t image_width, const uint32_t image_height, @@ -365,9 +365,9 @@ projection_2dgs_packed_bwd( 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 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] @@ -394,7 +394,7 @@ rasterize_to_pixels_2dgs_fwd( 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] + const at::optional masks, // [..., tile_height, tile_width] // image size const uint32_t image_width, const uint32_t image_height, @@ -420,7 +420,7 @@ rasterize_to_pixels_2dgs_bwd( 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] + const at::optional masks, // [..., tile_height, tile_width] // image size const uint32_t image_width, const uint32_t image_height, @@ -429,7 +429,8 @@ rasterize_to_pixels_2dgs_bwd( 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_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] @@ -475,8 +476,8 @@ projection_ut_3dgs_fused( 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] + 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, @@ -488,10 +489,12 @@ projection_ut_3dgs_fused( // uncented transform const UnscentedTransformParameters ut_params, ShutterType rs_type, - const at::optional radial_coeffs, // [..., C, 6] or [..., C, 4] optional + 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 + const at::optional thin_prism_coeffs, // [..., C, 4] optional + const FThetaCameraDistortionParameters + ftheta_coeffs // shared parameters for all cameras ); std::tuple @@ -503,24 +506,26 @@ rasterize_to_pixels_from_world_3dgs_fwd( 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] + 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::Tensor viewmats0, // [..., C, 4, 4] const at::optional - viewmats1, // [..., C, 4, 4] optional for rolling shutter - const at::Tensor Ks, // [..., C, 3, 3] + 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 + 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 + 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] @@ -535,30 +540,32 @@ rasterize_to_pixels_from_world_3dgs_bwd( 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] + 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::Tensor viewmats0, // [..., C, 4, 4] const at::optional - viewmats1, // [..., C, 4, 4] optional for rolling shutter - const at::Tensor Ks, // [..., C, 3, 3] + 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 + 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 + 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] + 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] + 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] diff --git a/gsplat/sycl/include/gsplat_sycl_utils.hpp b/gsplat/sycl/include/gsplat_sycl_utils.hpp index 76bf074f..2212587f 100644 --- a/gsplat/sycl/include/gsplat_sycl_utils.hpp +++ b/gsplat/sycl/include/gsplat_sycl_utils.hpp @@ -1,73 +1,71 @@ #ifndef GSPLAT_SYCL_UTILS #define GSPLAT_SYCL_UTILS +#include -#include - -template -struct BufferType { - using type = sycl::marray; +template struct BufferType { + using type = sycl::marray; constexpr static bool isVec{false}; }; -template -struct BufferType { - using type = sycl::vec; +template struct BufferType { + using type = sycl::vec; constexpr static bool isVec{true}; }; -template -struct BufferType { - using type = sycl::vec; +template struct BufferType { + using type = sycl::vec; constexpr static bool isVec{true}; }; -template -struct BufferType { - using type = sycl::vec; +template struct BufferType { + using type = sycl::vec; constexpr static bool isVec{true}; }; -template -struct BufferType { - using type = sycl::vec; +template struct BufferType { + using type = sycl::vec; constexpr static bool isVec{true}; }; -template -struct BufferType { - using type = sycl::vec; +template struct BufferType { + using type = sycl::vec; constexpr static bool isVec{true}; }; -template -using BufferType_t = typename BufferType::type; +template +using BufferType_t = typename BufferType::type; -template -void readToBuffer(T& dest, const void *source) { - dest = *(reinterpret_cast< const T *>(source)); +template void readToBuffer(T &dest, const void *source) { + dest = *(reinterpret_cast(source)); } -template -void gpuAtomicAdd(T* ptr, T value) { - sycl::atomic_ref +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 +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 +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); } diff --git a/gsplat/sycl/include/helpers.hpp b/gsplat/sycl/include/helpers.hpp index 4372d76d..5abd1c65 100644 --- a/gsplat/sycl/include/helpers.hpp +++ b/gsplat/sycl/include/helpers.hpp @@ -1,14 +1,16 @@ #ifndef GSPLAT_SYCL_HELPERS_HPP #define GSPLAT_SYCL_HELPERS_HPP -#include +#include -template -void gpuAtomicAdd(T* ptr, T value) { - sycl::atomic_ref +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); } -#endif //GSPLAT_SYCL_HELPERS_HPP \ No newline at end of file +#endif // GSPLAT_SYCL_HELPERS_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp index 1859e414..6b958954 100644 --- a/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp @@ -1,39 +1,37 @@ #ifndef ComputeShBwdKernel_HPP #define ComputeShBwdKernel_HPP -#include "utils.hpp" #include "spherical_harmonics.hpp" #include "types.hpp" +#include "utils.hpp" namespace gsplat::xpu { -template -struct ComputeShBwdKernel{ +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 + 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 + 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) - {} + : 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); @@ -56,14 +54,14 @@ struct ComputeShBwdKernel{ 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); + 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); } } }; -#endif //ComputeShBwdKernel_HPP +#endif // ComputeShBwdKernel_HPP } // 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 index 963326ba..93def949 100644 --- a/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp @@ -5,32 +5,28 @@ namespace gsplat::xpu { -template -struct ComputeShFwdKernel{ +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] + 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) - {} + 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 - { + void operator()(sycl::nd_item<1> work_item) const { uint32_t idx = work_item.get_global_id(0); if (idx >= m_N * 3) { return; @@ -50,6 +46,6 @@ struct ComputeShFwdKernel{ } }; -#endif //ComputeShFwdKernel_HPP +#endif // ComputeShFwdKernel_HPP } // 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 index 70e210e3..de5b29a8 100644 --- a/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp @@ -1,111 +1,114 @@ #ifndef FullyFusedProjectionBwdKernel_HPP #define FullyFusedProjectionBwdKernel_HPP - -#include "utils.hpp" +#include "proj.hpp" #include "quat.hpp" #include "quat_scale_to_covar_preci.hpp" -#include "proj.hpp" #include "transform.hpp" +#include "utils.hpp" namespace gsplat::xpu { -template -struct FullyFusedProjectionBwdKernel{ +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 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 + 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 + 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 + 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 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 + 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) - {} + // 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 - { + 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)) { + 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 + 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; + 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 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); @@ -137,7 +140,7 @@ struct FullyFusedProjectionBwdKernel{ vec3 scale; if (m_covars != nullptr) { // Changed: Updated pointer arithmetic - const T* covars = m_covars + bid * m_N * 6 + gid * 6; + const T *covars = m_covars + bid * m_N * 6 + gid * 6; covar = mat3( covars[0], covars[1], @@ -167,54 +170,54 @@ struct FullyFusedProjectionBwdKernel{ 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; + 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 @@ -232,8 +235,8 @@ struct FullyFusedProjectionBwdKernel{ if (m_v_means != nullptr) { // Changed: Updated pointer arithmetic - T* v_means = m_v_means + bid * m_N * 3 + gid * 3; - #pragma unroll + 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]); } @@ -241,7 +244,7 @@ struct FullyFusedProjectionBwdKernel{ if (m_v_covars != nullptr) { // Changed: Updated pointer arithmetic - T* v_covars = m_v_covars + bid * m_N * 6 + gid * 6; + 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]); @@ -257,8 +260,8 @@ struct FullyFusedProjectionBwdKernel{ 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; + 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]); @@ -270,10 +273,10 @@ struct FullyFusedProjectionBwdKernel{ if (m_v_viewmats != nullptr) { // Changed: Updated pointer arithmetic - T* v_viewmats = m_v_viewmats + bid * m_C * 16 + cid * 16; - #pragma unroll + 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 +#pragma unroll for (uint32_t j = 0; j < 3; j++) { // cols gpuAtomicAdd(v_viewmats + i * 4 + j, v_R[j][i]); } @@ -283,6 +286,6 @@ struct FullyFusedProjectionBwdKernel{ } }; -#endif //FullyFusedProjectionBwdKernel_HPP +#endif // FullyFusedProjectionBwdKernel_HPP } // 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 index e5074e46..5bae3b52 100644 --- a/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp @@ -1,28 +1,26 @@ #ifndef FullyFusedProjectionFwdKernel_HPP #define FullyFusedProjectionFwdKernel_HPP - -#include "utils.hpp" -#include "quat_scale_to_covar_preci.hpp" #include "proj.hpp" +#include "quat_scale_to_covar_preci.hpp" #include "transform.hpp" +#include "utils.hpp" namespace gsplat::xpu { -template -struct FullyFusedProjectionFwdKernel{ +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 + 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 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; @@ -32,25 +30,25 @@ struct FullyFusedProjectionFwdKernel{ 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 + 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, + 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 T *opacities, + const T *viewmats, + const T *Ks, const int32_t image_width, const int32_t image_height, const T eps2d, @@ -58,22 +56,23 @@ struct FullyFusedProjectionFwdKernel{ const T far_plane, const T radius_clip, const CameraModelType camera_model, - int32_t * radii, - T* means2d, - T* depths, - T* conics, - T* compensations + 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) - {} + // 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 - { + 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) { @@ -82,12 +81,12 @@ struct FullyFusedProjectionFwdKernel{ // 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 + 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 *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( @@ -117,7 +116,7 @@ struct FullyFusedProjectionFwdKernel{ mat3 covar; if (m_covars != nullptr) { // Changed: Updated pointer arithmetic - const T* covars = m_covars + bid * m_N * 6 + gid * 6; + const T *covars = m_covars + bid * m_N * 6 + gid * 6; covar = mat3( covars[0], covars[1], @@ -132,8 +131,8 @@ struct FullyFusedProjectionFwdKernel{ } 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; + 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 ); @@ -146,48 +145,48 @@ struct FullyFusedProjectionFwdKernel{ 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; + 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; @@ -216,7 +215,9 @@ struct FullyFusedProjectionFwdKernel{ m_radii[idx * 2 + 1] = 0; return; } - extend = sycl::min(extend, sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD))); + extend = sycl::min( + extend, sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD)) + ); } T radius_x = sycl::ceil(extend * sycl::sqrt(covar2d[0][0])); @@ -249,10 +250,8 @@ struct FullyFusedProjectionFwdKernel{ if (m_compensations != nullptr) { m_compensations[idx] = compensation; } - } - }; -#endif //FullyFusedProjectionFwdKernel_HPP +#endif // FullyFusedProjectionFwdKernel_HPP } // 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 index a6fba80c..b64f4f62 100644 --- a/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp +++ b/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp @@ -6,29 +6,24 @@ namespace gsplat::xpu { struct IsectOffsetEncodeKernel { const uint32_t m_n_isects; - const int64_t* m_isect_ids; + 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] + int32_t *m_offsets; //[C, n_tiles] IsectOffsetEncodeKernel( const uint32_t n_isects, - const int64_t* isect_ids, + 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) - {} + 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 { + void operator()(sycl::nd_item<1> work_item) const { uint32_t idx = work_item.get_global_id(0); if (idx >= m_n_isects) @@ -53,10 +48,11 @@ struct IsectOffsetEncodeKernel { 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 + 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); @@ -67,6 +63,6 @@ struct IsectOffsetEncodeKernel { } }; -#endif +#endif } // 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 index 69f03f3d..ee91751f 100644 --- a/gsplat/sycl/include/kernels/IsectTilesKernel.hpp +++ b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp @@ -1,9 +1,8 @@ #ifndef IsectTilesKernel_HPP #define IsectTilesKernel_HPP - -#include "types.hpp" #include "transform.hpp" +#include "types.hpp" #include "utils.hpp" #include @@ -14,65 +13,53 @@ struct uint2 { uint32_t y; }; -template -struct IsectTilesKernel { +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 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] + 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 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 { + 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; @@ -98,10 +85,30 @@ struct IsectTilesKernel { 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); + 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 @@ -124,9 +131,9 @@ struct IsectTilesKernel { const int64_t cid_enc = cid << (32 + m_tile_n_bits); - int32_t depth_i32 = *reinterpret_cast(&m_depths[idx]); + 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) { @@ -142,6 +149,6 @@ struct IsectTilesKernel { } }; -#endif //IsectTilesKernel_HPP +#endif // IsectTilesKernel_HPP } // 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 index 3bf876bb..53087853 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp @@ -1,69 +1,91 @@ #ifndef PackedProjectionBwdKernel_HPP #define PackedProjectionBwdKernel_HPP -#include -#include "utils.hpp" +#include "proj.hpp" #include "quat.hpp" #include "quat_scale_to_covar_preci.hpp" -#include "proj.hpp" #include "transform.hpp" +#include "utils.hpp" +#include namespace gsplat::xpu { -template -struct PackedProjectionBwdKernel { +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 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; + 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 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; + 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, + 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) - {} + 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); @@ -79,39 +101,60 @@ struct PackedProjectionBwdKernel { 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]); + 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); + 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; + + 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] + 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; + 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] + 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); @@ -126,18 +169,57 @@ struct PackedProjectionBwdKernel { mat3 v_covar_c(0.f); vec3 v_mean_c(0.f); - const T* v_means2d = m_v_means2d + idx * 2; - + 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; + 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]; @@ -146,7 +228,9 @@ struct PackedProjectionBwdKernel { 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); + 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 --- @@ -154,11 +238,13 @@ struct PackedProjectionBwdKernel { 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; + 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; + 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]; @@ -169,50 +255,105 @@ struct PackedProjectionBwdKernel { 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; + 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; + T *v_means_out = m_v_means + bid * m_N * 3 + gid * 3; for (int i = 0; i < 3; ++i) { - sycl::atomic_ref ref(v_means_out[i]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device> + ref(v_means_out[i]); ref.fetch_add(v_mean[i]); } } if (m_v_covars != nullptr) { - T* v_covars_out = m_v_covars + bid * m_N * 6 + gid * 6; - sycl::atomic_ref(v_covars_out[0]).fetch_add(v_covar[0][0]); - sycl::atomic_ref(v_covars_out[1]).fetch_add(v_covar[0][1] + v_covar[1][0]); - sycl::atomic_ref(v_covars_out[2]).fetch_add(v_covar[0][2] + v_covar[2][0]); - sycl::atomic_ref(v_covars_out[3]).fetch_add(v_covar[1][1]); - sycl::atomic_ref(v_covars_out[4]).fetch_add(v_covar[1][2] + v_covar[2][1]); - sycl::atomic_ref(v_covars_out[5]).fetch_add(v_covar[2][2]); + T *v_covars_out = m_v_covars + bid * m_N * 6 + gid * 6; + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_covars_out[0]) + .fetch_add(v_covar[0][0]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_covars_out[1]) + .fetch_add(v_covar[0][1] + v_covar[1][0]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_covars_out[2]) + .fetch_add(v_covar[0][2] + v_covar[2][0]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_covars_out[3]) + .fetch_add(v_covar[1][1]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_covars_out[4]) + .fetch_add(v_covar[1][2] + v_covar[2][1]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_covars_out[5]) + .fetch_add(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) sycl::atomic_ref(v_quats_out[i]).fetch_add(v_quat[i]); - for (int i = 0; i < 3; ++i) sycl::atomic_ref(v_scales_out[i]).fetch_add(v_scale[i]); + 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) + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_quats_out[i]) + .fetch_add(v_quat[i]); + for (int i = 0; i < 3; ++i) + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device>(v_scales_out[i]) + .fetch_add(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 + 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 - sycl::atomic_ref ref(v_viewmats_out[i * 4 + j]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device> + ref(v_viewmats_out[i * 4 + j]); ref.fetch_add(v_R[j][i]); } - sycl::atomic_ref ref(v_viewmats_out[i * 4 + 3]); + sycl::atomic_ref< + T, + sycl::memory_order::relaxed, + sycl::memory_scope::device> + ref(v_viewmats_out[i * 4 + 3]); ref.fetch_add(v_t[i]); } } diff --git a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp index 36f4b328..66b51867 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -1,27 +1,26 @@ #ifndef PackedProjectionFwdKernel_HPP #define PackedProjectionFwdKernel_HPP -#include -#include "utils.hpp" -#include "quat_scale_to_covar_preci.hpp" #include "proj.hpp" +#include "quat_scale_to_covar_preci.hpp" #include "transform.hpp" +#include "utils.hpp" +#include namespace gsplat::xpu { -template -struct PackedProjectionFwdKernel { +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 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; @@ -29,57 +28,79 @@ struct PackedProjectionFwdKernel { const T m_far_plane; const T m_radius_clip; const CameraModelType m_camera_model; - const int32_t* m_block_accum; // Packing helper for the second pass + const int32_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; + 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, + 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 int32_t* block_accum, + const int32_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) - {} + 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_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 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 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; @@ -92,16 +113,24 @@ struct PackedProjectionFwdKernel { 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; - + 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] + 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] ); - 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; @@ -115,31 +144,75 @@ struct PackedProjectionFwdKernel { if (valid) { mat3 covar; if (m_covars != nullptr) { - const T* current_covars = m_covars + bid * m_N * 6 + gid * 6; + 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] + 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); + 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; + 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; + 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); @@ -149,7 +222,7 @@ struct PackedProjectionFwdKernel { inverse(covar2d, covar2d_inv); } } - + T radius_x, radius_y; if (valid) { const T ALPHA_THRESHOLD = 1.f / 255.f; @@ -162,7 +235,10 @@ struct PackedProjectionFwdKernel { if (opacity < ALPHA_THRESHOLD) { valid = false; } - extend = sycl::fmin(extend, sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD))); + extend = sycl::fmin( + extend, + sycl::sqrt(2.0f * sycl::log(opacity / ALPHA_THRESHOLD)) + ); } radius_x = sycl::ceil(extend * sycl::sqrt(covar2d[0][0])); @@ -172,8 +248,10 @@ struct PackedProjectionFwdKernel { 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) { + 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; } } @@ -186,12 +264,13 @@ struct PackedProjectionFwdKernel { 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<>()); + 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) { + if (local_id == 0) { m_block_cnts[block_idx] = 0; } } @@ -200,15 +279,18 @@ struct PackedProjectionFwdKernel { // Second pass: Write data for visible Gaussians. bool any_valid = sycl::any_of_group(group, valid); if (any_valid) { - // Perform an exclusive scan to find the local offset for this thread. - int32_t local_offset = sycl::exclusive_scan_over_group(group, thread_data, sycl::plus<>()); - + // Perform an exclusive scan to find the local offset for this + // thread. + int32_t local_offset = sycl::exclusive_scan_over_group( + group, thread_data, sycl::plus<>() + ); + if (valid) { int32_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; @@ -230,8 +312,9 @@ struct PackedProjectionFwdKernel { 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]; + // 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]; } diff --git a/gsplat/sycl/include/kernels/ProjBwdKernel.hpp b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp index 9d363fa6..b75a6eb7 100644 --- a/gsplat/sycl/include/kernels/ProjBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp @@ -1,67 +1,65 @@ #ifndef ProjBwdKernel_HPP #define ProjBwdKernel_HPP - -#include "proj.hpp" #include "Common.h" +#include "proj.hpp" namespace gsplat::xpu { -template -struct ProjBwdKernel{ +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 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] + 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 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) - {} + 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)); + 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); + 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); + 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); @@ -72,71 +70,71 @@ struct ProjBwdKernel{ 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; + 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 +// 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 +#pragma unroll for (uint32_t j = 0; j < 3; j++) { // cols v_covars[i * 3 + j] = T(v_covar[j][i]); } } - #pragma unroll +#pragma unroll for (uint32_t i = 0; i < 3; i++) { v_means[i] = T(v_mean[i]); } } }; -#endif //ProjBwdKernel_HPP +#endif // ProjBwdKernel_HPP } // 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 index 8cef2b28..d7b28ae4 100644 --- a/gsplat/sycl/include/kernels/ProjFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp @@ -1,60 +1,58 @@ #ifndef ProjFwdKernel_HPP #define ProjFwdKernel_HPP - -#include "proj.hpp" #include "Common.h" +#include "proj.hpp" namespace gsplat::xpu { -template -struct ProjFwdKernel{ +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 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] + 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 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] + 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) - {} + : 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)); + 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); - + 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); @@ -62,31 +60,36 @@ struct ProjFwdKernel{ 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; + 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 +#pragma unroll for (uint32_t i = 0; i < 2; i++) { // rows - #pragma unroll +#pragma unroll for (uint32_t j = 0; j < 2; j++) { // cols covars2d[i * 2 + j] = T(covar2d[j][i]); } } - #pragma unroll +#pragma unroll for (uint32_t i = 0; i < 2; i++) { means2d[i] = T(mean2d[i]); } } - }; -#endif //ProjFwdKernel_HPP +#endif // ProjFwdKernel_HPP -} //namespace gsplat::xpu \ No newline at end of file +} // 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 index 7272e6ae..4a05409f 100644 --- a/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp @@ -1,16 +1,15 @@ #ifndef Projection2DGSFusedBwdKernel_HPP #define Projection2DGSFusedBwdKernel_HPP -#include "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 T sum(vec3 a) { return a.x + a.y + a.z; } -template +template inline void compute_ray_transforms_aabb_vjp( const T *ray_transforms, const T *v_means2d, @@ -31,8 +30,8 @@ inline void compute_ray_transforms_aabb_vjp( ) { 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]; + 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]; @@ -40,14 +39,26 @@ inline void compute_ray_transforms_aabb_vjp( 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]; + 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; @@ -85,75 +96,82 @@ inline void compute_ray_transforms_aabb_vjp( 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 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 { +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 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] + 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] + 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] + 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 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 + 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 - { + : 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; } @@ -162,51 +180,62 @@ struct Projection2DGSFusedBwdKernel { 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 *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 *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; + 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 + 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]); + 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) + 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] + 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 @@ -242,27 +271,27 @@ struct Projection2DGSFusedBwdKernel { // Write out results with atomic additions if (m_v_means != nullptr) { - T* v_means_out = m_v_means + bid * m_N * 3 + gid * 3; + 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; - + 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; - + 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++) { diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp index 8afa122f..cf9b28fb 100644 --- a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp @@ -1,89 +1,93 @@ #ifndef Projection2DGSFusedFwdKernel_HPP #define Projection2DGSFusedFwdKernel_HPP -#include "utils.hpp" #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 inline float sum(vec3 a) { return a.x + a.y + a.z; } -template -struct Projection2DGSFusedFwdKernel { +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 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] + 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 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 + 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 - { + : 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; + 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 + 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]); @@ -98,42 +102,48 @@ struct Projection2DGSFusedFwdKernel { return; } - const T* quats = m_quats + bid * m_N * 4 + gid * 4; - const T* scales = m_scales + bid * m_N * 3 + gid * 3; + 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])); + 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) + 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 - ); + 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) + 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 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)); @@ -153,8 +163,10 @@ struct Projection2DGSFusedFwdKernel { 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))); + 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; @@ -163,15 +175,18 @@ struct Projection2DGSFusedFwdKernel { } // 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) { + 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 = 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 @@ -186,16 +201,15 @@ struct Projection2DGSFusedFwdKernel { 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] - + 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; diff --git a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp index 1baf9400..667fc721 100644 --- a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp @@ -1,54 +1,51 @@ #ifndef QuatScaleToCovarPreciBwdKernel_HPP #define QuatScaleToCovarPreciBwdKernel_HPP - #include "quat_scale_to_covar_preci.hpp" namespace gsplat::xpu { - -template -struct QuatScaleToCovarPreciBwdKernel{ + +template struct QuatScaleToCovarPreciBwdKernel { const uint32_t m_N; // fwd inputs - const T* m_quats; // [N, 4] - const T* m_scales; // [N, 3] + 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 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] + 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 T *quats, + const T *scales, + const T *v_covars, + const T *v_precis, const bool triu, - T* v_scales, - T* v_quats + 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) - {} + : 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 - { + 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); + 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); @@ -56,7 +53,7 @@ struct QuatScaleToCovarPreciBwdKernel{ // glm is column-major, input is row-major mat3 v_covar; if (m_triu) { - const T* v_covars = m_v_covars + (idx * 6); + const T *v_covars = m_v_covars + (idx * 6); v_covar = mat3( v_covars[0], v_covars[1] * .5f, @@ -69,7 +66,7 @@ struct QuatScaleToCovarPreciBwdKernel{ v_covars[5] ); } else { - const T* v_covars = m_v_covars + (idx * 9); + 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); } @@ -82,7 +79,7 @@ struct QuatScaleToCovarPreciBwdKernel{ // glm is column-major, input is row-major mat3 v_preci; if (m_triu) { - const T* v_precis = m_v_precis + (idx * 6); + const T *v_precis = m_v_precis + (idx * 6); v_preci = mat3( v_precis[0], v_precis[1] * .5f, @@ -95,7 +92,7 @@ struct QuatScaleToCovarPreciBwdKernel{ v_precis[5] ); } else { - const T* v_precis = m_v_precis + (idx * 9); + 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); } @@ -104,18 +101,17 @@ struct QuatScaleToCovarPreciBwdKernel{ ); } - #pragma unroll +#pragma unroll for (uint32_t k = 0; k < 3; ++k) { v_scales[k] = T(v_scale[k]); } - #pragma unroll +#pragma unroll for (uint32_t k = 0; k < 4; ++k) { v_quats[k] = T(v_quat[k]); } - } - + } }; -#endif //QuatScaleToCovarPreciBwdKernel_HPP +#endif // QuatScaleToCovarPreciBwdKernel_HPP } // 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 index a82961f1..8e587b25 100644 --- a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp @@ -1,33 +1,30 @@ #ifndef QuatScaleToCovarPreciFwdKernel_HPP #define QuatScaleToCovarPreciFwdKernel_HPP - #include "quat_scale_to_covar_preci.hpp" namespace gsplat::xpu { - -template -struct QuatScaleToCovarPreciFwdKernel{ + +template struct QuatScaleToCovarPreciFwdKernel { const uint32_t m_N; - const T* m_quats; // [N, 4] - const T* m_scales; // [N, 3] + 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] + 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 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) - {} + 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); @@ -35,20 +32,23 @@ struct QuatScaleToCovarPreciFwdKernel{ return; } - const T* quats = m_quats + (idx * 4); - const T* scales = m_scales + (idx * 3); + 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 + 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); + T *covars = m_covars + (idx * 6); covars[0] = T(covar[0][0]); covars[1] = T(covar[0][1]); covars[2] = T(covar[0][2]); @@ -56,10 +56,10 @@ struct QuatScaleToCovarPreciFwdKernel{ covars[4] = T(covar[1][2]); covars[5] = T(covar[2][2]); } else { - T* covars = m_covars + (idx * 9); - #pragma unroll + T *covars = m_covars + (idx * 9); +#pragma unroll for (uint32_t i = 0; i < 3; i++) { // rows - #pragma unroll +#pragma unroll for (uint32_t j = 0; j < 3; j++) { // cols covars[i * 3 + j] = T(covar[j][i]); } @@ -69,7 +69,7 @@ struct QuatScaleToCovarPreciFwdKernel{ if (m_precis != nullptr) { if (m_triu) { - T* precis = m_precis + (idx * 6); + T *precis = m_precis + (idx * 6); precis[0] = T(preci[0][0]); precis[1] = T(preci[0][1]); precis[2] = T(preci[0][2]); @@ -77,18 +77,18 @@ struct QuatScaleToCovarPreciFwdKernel{ precis[4] = T(preci[1][2]); precis[5] = T(preci[2][2]); } else { - T* precis = m_precis + (idx * 9); - #pragma unroll + T *precis = m_precis + (idx * 9); +#pragma unroll for (uint32_t i = 0; i < 3; i++) { // rows - #pragma unroll +#pragma unroll for (uint32_t j = 0; j < 3; j++) { // cols precis[i * 3 + j] = T(preci[j][i]); } } } } - } + } }; -#endif //QuatScaleToCovarPreciFwdKernel_HPP +#endif // QuatScaleToCovarPreciFwdKernel_HPP } // 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 index 6ceffb12..fdf2dbab 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp @@ -1,9 +1,9 @@ #ifndef RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP #define RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP -#include -#include "types.hpp" #include "gsplat_sycl_utils.hpp" +#include "types.hpp" +#include namespace gsplat::xpu { @@ -11,57 +11,57 @@ namespace gsplat::xpu { constexpr float ALPHA_THRESHOLD = 1.0f / 255.0f; constexpr float FILTER_INV_SQUARE_2DGS = 2.0f; -template -struct RasterizeToPixels2DGSBwdKernel { +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 - + 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 - + 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 - + 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 - + 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 - + 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; @@ -78,13 +78,13 @@ struct RasterizeToPixels2DGSBwdKernel { 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, + 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, @@ -92,27 +92,27 @@ struct RasterizeToPixels2DGSBwdKernel { const uint32_t tile_width, const uint32_t tile_height, // Intersections - const int32_t* tile_offsets, - const int32_t* flatten_ids, + 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, + 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, + 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, + 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, @@ -121,172 +121,197 @@ struct RasterizeToPixels2DGSBwdKernel { 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) - {} + ) + : 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 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 - + + 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; + 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; + 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; - + 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; + const float *backgrounds_ptr = m_backgrounds; if (backgrounds_ptr != nullptr) { backgrounds_ptr += image_id * COLOR_DIM; } - - const bool* masks_ptr = m_masks; + + 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)) - ); - + 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))) { + 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; - + 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); + 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; + 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_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() + 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); - + 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], @@ -303,102 +328,117 @@ struct RasterizeToPixels2DGSBwdKernel { 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*>(m_colors + g * COLOR_DIM); + 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] + 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) { + 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; + 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}; - + 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]); - + 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]); - + 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); - + 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); + 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}; @@ -408,43 +448,46 @@ struct RasterizeToPixels2DGSBwdKernel { 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]; + 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]; + 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; @@ -453,143 +496,217 @@ struct RasterizeToPixels2DGSBwdKernel { } 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; + 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; + 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 + + // 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]) + 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); - + 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_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] + 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 + + // 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); - + 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]) + 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>()); + 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()); + 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()); - + 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>()); + 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]); + 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]); + 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]); - + 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]); + 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; + m_v_densify[g * 2 + 1] = + m_v_ray_transforms[g * 9 + 5] * depth; } } } diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp index 6bcd5b56..06105089 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp @@ -1,9 +1,9 @@ #ifndef RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP #define RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP -#include -#include "types.hpp" #include "gsplat_sycl_utils.hpp" +#include "types.hpp" +#include namespace gsplat::xpu { @@ -11,188 +11,200 @@ namespace gsplat::xpu { 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 +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 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 - + + 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 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, + 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) - {} + ) + : 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 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 - + + 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; - + 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; + const float *backgrounds_ptr = m_backgrounds; if (backgrounds_ptr != nullptr) { backgrounds_ptr += image_id * COLOR_DIM; } - - const bool* masks_ptr = m_masks; + + 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] = + 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)) + 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; - + 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 - + 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], @@ -210,25 +222,27 @@ struct RasterizeToPixels2DGSFwdKernel { 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 - + + // 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); + 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( @@ -236,14 +250,14 @@ struct RasterizeToPixels2DGSFwdKernel { 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( @@ -251,71 +265,74 @@ struct RasterizeToPixels2DGSFwdKernel { 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] + 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]); - + 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); - + 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*>(m_colors + g * COLOR_DIM); + 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; + 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]; @@ -324,28 +341,31 @@ struct RasterizeToPixels2DGSFwdKernel { 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; + 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]; @@ -353,33 +373,36 @@ struct RasterizeToPixels2DGSFwdKernel { } } else { // With background - if constexpr(BufferType::isVec && COLOR_DIM <= 4) { + 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; + *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]; + 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); diff --git a/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp index 75c237ee..f0ff4fbc 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp @@ -1,368 +1,448 @@ #ifndef RasterizeToPixelsBwdKernel_HPP #define RasterizeToPixelsBwdKernel_HPP - -#include "types.hpp" #include "gsplat_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]; +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))) - { + if ((camera_id == m_C - 1) && + (tile_id == static_cast(m_tile_width * m_tile_height - 1) + )) { range_end = m_n_isects; - } - else - { + } 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; + 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]; + 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(); - 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(); - - auto conic = *(reinterpret_cast*>(data+2) ); - m_slm_conics[threadRank] = conic.template convert(); - - if constexpr(BufferType::isVec && COLOR_DIM <= 4){ - auto color = *( reinterpret_cast*>(data + 2 + 3) ); - m_slm_colors[threadRank] = color.template convert();; - } - } - m_slm_opacities[threadRank] = static_cast(*(data + 2 + 3 + COLOR_DIM)); - - } else { - m_slm_means2d[threadRank] = m_means2d[g].template convert(); - - m_slm_opacities[threadRank] = static_cast(m_opacities[g]); - auto temp = *( reinterpret_cast*>(m_conics + g) ); - - m_slm_conics[threadRank] = temp.template convert(); - - if constexpr(BufferType::isVec && COLOR_DIM <= 4){ - auto temp2 = *( reinterpret_cast*>(m_colors + g * COLOR_DIM) ); - m_slm_colors[threadRank] = temp2.template convert(); - } - } - } - - 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(); - } 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]); - } - } - } - } - } + 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] + ); + } + } + } + } + } }; #endif // RasterizeToPixelsBwdKernel_HPP diff --git a/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp index a25718a8..8528c1fa 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp @@ -1,269 +1,325 @@ #ifndef RasterizeToPixelsFwdKernel_HPP #define RasterizeToPixelsFwdKernel_HPP -#include "types.hpp" #include "gsplat_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*>(data + 2 + 3) ); - m_slm_opacities[threadRank] = *(data + 2 + 3 + COLOR_DIM); - } - } +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 { - 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) ); - } + 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); + } } - } - - 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; + } 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) + ); } - - - 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]; - } - } - } + } + } + + 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]; + } + } + } }; -#endif //RasterizeToPixelsFwdKernel_HPP +#endif // RasterizeToPixelsFwdKernel_HPP } // 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 index 29c544de..9fc88eb0 100644 --- a/gsplat/sycl/include/kernels/RelocationKernel.hpp +++ b/gsplat/sycl/include/kernels/RelocationKernel.hpp @@ -42,8 +42,8 @@ template class RelocationKernel { 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(-1.0f, k) / sycl::sqrt(static_cast(k + 1)) + ) * sycl::pow(static_cast(new_opacities[idx]), k + 1); denom_sum += (bin_coeff * term); } diff --git a/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp index 004af61f..0ef6613c 100644 --- a/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp @@ -1,45 +1,42 @@ #ifndef WorldToCamBwdKernel_HPP #define WorldToCamBwdKernel_HPP - -#include "types.hpp" #include "transform.hpp" +#include "types.hpp" #include "utils.hpp" namespace gsplat::xpu { - -template -struct WorldToCamBwdKernel{ + +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] + 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 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) { @@ -50,9 +47,9 @@ struct WorldToCamBwdKernel{ 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); + 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( @@ -68,7 +65,7 @@ struct WorldToCamBwdKernel{ ); 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); @@ -80,36 +77,37 @@ struct WorldToCamBwdKernel{ 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_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 + 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]); - } + gpuAtomicAdd(v_means + i, v_mean[i]); + } } if (m_v_covars != nullptr) { - T* v_covars = m_v_covars + (gid * 9); - #pragma unroll + T *v_covars = m_v_covars + (gid * 9); +#pragma unroll for (uint32_t i = 0; i < 3; i++) { // rows - #pragma unroll +#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 + T *v_viewmats = m_v_viewmats + cid * 16; +#pragma unroll for (uint32_t i = 0; i < 3; i++) { // rows - #pragma unroll +#pragma unroll for (uint32_t j = 0; j < 3; j++) { // cols gpuAtomicAdd(v_viewmats + i * 4 + j, v_R[j][i]); } @@ -119,6 +117,6 @@ struct WorldToCamBwdKernel{ } }; -#endif //WorldToCamBwdKernel_HPP +#endif // WorldToCamBwdKernel_HPP } // 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 index 9296a715..59829a48 100644 --- a/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp @@ -3,40 +3,37 @@ /**************************************************************************** * World to Camera Transformation Forward Pass - * From: https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/csrc/world_to_cam_fwd.cu + * From: + *https://github.com/nerfstudio-project/gsplat/blob/main/gsplat/cuda/csrc/world_to_cam_fwd.cu ****************************************************************************/ -#include "types.hpp" #include "transform.hpp" +#include "types.hpp" namespace gsplat::xpu { - -template -struct WorldToCamFwdKernel{ + +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] + 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 + 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) - {} + : 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 - { + 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; @@ -46,10 +43,10 @@ struct WorldToCamFwdKernel{ 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); - + 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], @@ -69,8 +66,8 @@ struct WorldToCamFwdKernel{ 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 + 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]; } @@ -81,10 +78,10 @@ struct WorldToCamFwdKernel{ 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 + T *covars_c = m_covars_c + (idx * 9); +#pragma unroll for (uint32_t i = 0; i < 3; i++) { // rows - #pragma unroll +#pragma unroll for (uint32_t j = 0; j < 3; j++) { // cols covars_c[i * 3 + j] = T(covar_c[j][i]); } @@ -93,6 +90,6 @@ struct WorldToCamFwdKernel{ } }; -#endif //WorldToCamFwdKernel_HPP +#endif // WorldToCamFwdKernel_HPP } // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/proj.hpp b/gsplat/sycl/include/proj.hpp index a5e5505c..60a6de9c 100644 --- a/gsplat/sycl/include/proj.hpp +++ b/gsplat/sycl/include/proj.hpp @@ -3,7 +3,6 @@ #include "types.hpp" - template inline void ortho_proj( // inputs @@ -19,7 +18,7 @@ inline void ortho_proj( mat2 &cov2d, vec2 &mean2d ) { - T x = mean3d[0], y = mean3d[1];// z = mean3d[2]; + T x = mean3d[0], y = mean3d[1]; // z = mean3d[2]; // mat3x2 is 3 columns x 2 rows. mat3x2 J = mat3x2( @@ -341,5 +340,4 @@ inline void fisheye_proj_vjp( v_mean3d.z += dL_dtz_raw; } - #endif // GSPLAT_SYCL_PROJ_HPP diff --git a/gsplat/sycl/include/quat.hpp b/gsplat/sycl/include/quat.hpp index 049b8efc..2c125814 100644 --- a/gsplat/sycl/include/quat.hpp +++ b/gsplat/sycl/include/quat.hpp @@ -2,10 +2,8 @@ #define GSPLAT_SYCL_QUAT_HPP #include "types.hpp" - -template -inline mat3 quat_to_rotmat(const vec4 quat) { +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); @@ -55,5 +53,4 @@ quat_to_rotmat_vjp(const vec4 quat, const mat3 v_R, vec4 &v_quat) { v_quat += (v_quat_n - glm::dot(v_quat_n, quat_n) * quat_n) * inv_norm; } - #endif // GSPLAT_SYCL_QUAT_HPP diff --git a/gsplat/sycl/include/quat_scale_to_covar_preci.hpp b/gsplat/sycl/include/quat_scale_to_covar_preci.hpp index 48dc2e07..f1f43a74 100644 --- a/gsplat/sycl/include/quat_scale_to_covar_preci.hpp +++ b/gsplat/sycl/include/quat_scale_to_covar_preci.hpp @@ -1,9 +1,8 @@ #ifndef GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP #define GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP -#include "types.hpp" #include "quat.hpp" - +#include "types.hpp" template inline void quat_scale_to_covar_preci( diff --git a/gsplat/sycl/include/spherical_harmonics.hpp b/gsplat/sycl/include/spherical_harmonics.hpp index 45bcfdbb..a94ca68b 100644 --- a/gsplat/sycl/include/spherical_harmonics.hpp +++ b/gsplat/sycl/include/spherical_harmonics.hpp @@ -2,7 +2,6 @@ #define GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP #include "types.hpp" - // Evaluate spherical harmonics bases at unit direction for high orders using // approach described by Efficient Spherical Harmonic Evaluation, Peter-Pike @@ -358,5 +357,4 @@ inline void sh_coeffs_to_color_fast_vjp( } } - #endif // GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/types.hpp b/gsplat/sycl/include/types.hpp index 8048bd83..70f18a4e 100644 --- a/gsplat/sycl/include/types.hpp +++ b/gsplat/sycl/include/types.hpp @@ -3,7 +3,6 @@ #include - template using vec2 = glm::vec<2, T>; template using vec3 = glm::vec<3, T>; @@ -18,5 +17,4 @@ template using mat4 = glm::mat<4, 4, T>; template using mat3x2 = glm::mat<3, 2, T>; - #endif // GSPLAT_SYCL_TYPES_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/utils.hpp b/gsplat/sycl/include/utils.hpp index 094fb354..308537a0 100644 --- a/gsplat/sycl/include/utils.hpp +++ b/gsplat/sycl/include/utils.hpp @@ -3,18 +3,19 @@ #include "types.hpp" -#include +#include -template -void gpuAtomicAdd(T* ptr, T value) { - sycl::atomic_ref +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 -inline T inverse(const mat2 M, mat2 &Minv) { +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; @@ -40,7 +41,8 @@ inline T add_blur(const T eps2d, mat2 &covar, T &compensation) { 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)); + compensation = + sycl::sqrt(sycl::max(static_cast(0), det_orig / det_blur)); return det_blur; } @@ -80,5 +82,4 @@ inline void add_blur_vjp( eps2d * det_conic_blur); } - #endif // GSPLAT_SYCL_UTILS_HPP diff --git a/gsplat/sycl/src/adam.cpp b/gsplat/sycl/src/adam.cpp index fd030547..7abfd6d8 100644 --- a/gsplat/sycl/src/adam.cpp +++ b/gsplat/sycl/src/adam.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - void adam( at::Tensor ¶m, // [..., D] const at::Tensor ¶m_grad, // [..., D] diff --git a/gsplat/sycl/src/intersect_offset.cpp b/gsplat/sycl/src/intersect_offset.cpp index 3a50e3b0..3d5dc95b 100644 --- a/gsplat/sycl/src/intersect_offset.cpp +++ b/gsplat/sycl/src/intersect_offset.cpp @@ -1,9 +1,9 @@ #include -#include +#include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/IsectOffsetEncodeKernel.hpp" namespace gsplat::xpu { @@ -26,32 +26,30 @@ at::Tensor intersect_offset( 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(); + auto &d_queue = at::xpu::getCurrentXPUStream().queue(); - size_t numWorkGrps = (n_isects + GSPLAT_N_THREADS - 1) / GSPLAT_N_THREADS; + 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); - } - ); + + 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; } diff --git a/gsplat/sycl/src/intersect_tile.cpp b/gsplat/sycl/src/intersect_tile.cpp index 0c0e051f..5fe318b3 100644 --- a/gsplat/sycl/src/intersect_tile.cpp +++ b/gsplat/sycl/src/intersect_tile.cpp @@ -1,19 +1,20 @@ #include + #include -#include "Ops.h" #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::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 I, // -> maps to C const uint32_t tile_size, const uint32_t tile_width, const uint32_t tile_height, @@ -23,8 +24,10 @@ std::tuple intersect_tile( CHECK_CONTIGUOUS(means2d); CHECK_CONTIGUOUS(radii); CHECK_CONTIGUOUS(depths); - if (image_ids.has_value()) CHECK_CONTIGUOUS(image_ids.value()); - if (gaussian_ids.has_value()) CHECK_CONTIGUOUS(gaussian_ids.value()); + if (image_ids.has_value()) + CHECK_CONTIGUOUS(image_ids.value()); + if (gaussian_ids.has_value()) + CHECK_CONTIGUOUS(gaussian_ids.value()); const bool packed = segmented; const uint32_t C = I; @@ -35,8 +38,11 @@ std::tuple intersect_tile( 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."); + 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; @@ -50,28 +56,39 @@ std::tuple intersect_tile( ); } auto options = depths.options(); - at::Tensor tiles_per_gauss = at::empty_like(depths, options.dtype(at::kInt)); + 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."); + 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; + 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) { + auto e1 = d_queue.submit([&](sycl::handler &cgh) { IsectTilesKernel kernel( - packed, C, N, nnz, + 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, + tile_size, + tile_width, + tile_height, + tile_n_bits, tiles_per_gauss.data_ptr(), nullptr, // isect_ids nullptr // flatten_ids @@ -80,7 +97,8 @@ std::tuple intersect_tile( }); e1.wait(); - at::Tensor cum_tiles_per_gauss = at::cumsum(tiles_per_gauss.view({-1}), 0, at::kLong); + 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(); @@ -88,18 +106,24 @@ std::tuple intersect_tile( 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) { + auto e2 = d_queue.submit([&](sycl::handler &cgh) { IsectTilesKernel kernel( - packed, C, N, nnz, + 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, + tile_size, + tile_width, + tile_height, + tile_n_bits, nullptr, // tiles_per_gauss isect_ids.data_ptr(), flatten_ids.data_ptr() diff --git a/gsplat/sycl/src/null.cpp b/gsplat/sycl/src/null.cpp index 88310566..bfcecb03 100644 --- a/gsplat/sycl/src/null.cpp +++ b/gsplat/sycl/src/null.cpp @@ -1,13 +1,13 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" -namespace gsplat::xpu { +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 +} // 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 index 2b08b134..3756ffc7 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp @@ -1,11 +1,11 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/Projection2DGSFusedBwdKernel.hpp" namespace gsplat::xpu { - + std::tuple projection_2dgs_fused_bwd( // fwd inputs @@ -38,10 +38,19 @@ projection_2dgs_fused_bwd( CHECK_CONTIGUOUS(v_normals); CHECK_CONTIGUOUS(v_ray_transforms); - 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( + 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 @@ -54,22 +63,26 @@ projection_2dgs_fused_bwd( 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(); + 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; + 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) { + means.scalar_type(), + "projection_2dgs_fused_bwd", + [&] { + auto e = d_queue.submit([&](sycl::handler &cgh) { Projection2DGSFusedBwdKernel kernel( B, C, @@ -90,12 +103,16 @@ projection_2dgs_fused_bwd( v_means.data_ptr(), v_quats.data_ptr(), v_scales.data_ptr(), - viewmats_requires_grad ? v_viewmats.data_ptr() : nullptr + viewmats_requires_grad ? v_viewmats.data_ptr() + : nullptr + ); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel ); - 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); } diff --git a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp index 731e218b..b39f4e40 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp @@ -1,11 +1,11 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/Projection2DGSFusedFwdKernel.hpp" namespace gsplat::xpu { - + std::tuple< at::Tensor, at::Tensor, @@ -31,11 +31,22 @@ projection_2dgs_fused_fwd( CHECK_CONTIGUOUS(viewmats); CHECK_CONTIGUOUS(Ks); - 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]"); + 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 @@ -57,7 +68,7 @@ projection_2dgs_fused_fwd( at::DimVector out_shape_cn3 = batch_dims; out_shape_cn3.insert(out_shape_cn3.end(), {C, N, 3}); - // Output shape: [..., C, N, 3, 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}); @@ -72,15 +83,18 @@ projection_2dgs_fused_fwd( 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; + 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) { + means.scalar_type(), + "projection_2dgs_fused_fwd", + [&] { + auto e = d_queue.submit([&](sycl::handler &cgh) { Projection2DGSFusedFwdKernel kernel( B, C, @@ -101,10 +115,13 @@ projection_2dgs_fused_fwd( ray_transforms.data_ptr(), normals.data_ptr() ); - cgh.parallel_for(sycl::nd_range<1>(global_range, local_range), kernel); + cgh.parallel_for( + sycl::nd_range<1>(global_range, local_range), kernel + ); }); e.wait(); - }); + } + ); return std::make_tuple(radii, means2d, depths, ray_transforms, normals); } diff --git a/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp b/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp index cb20c2a3..bf36b64f 100644 --- a/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_packed_bwd.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple projection_2dgs_packed_bwd( // fwd inputs @@ -17,9 +17,9 @@ projection_2dgs_packed_bwd( 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 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] diff --git a/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp index 98cc43f9..45c0005d 100644 --- a/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_packed_fwd.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple< at::Tensor, at::Tensor, diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp index 6f0be199..b20314d9 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp @@ -1,10 +1,10 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/FullyFusedProjectionBwdKernel.hpp" -namespace gsplat::xpu { +namespace gsplat::xpu { std::tuple projection_ewa_3dgs_fused_bwd( @@ -32,18 +32,23 @@ projection_ewa_3dgs_fused_bwd( ) { // Input validation CHECK_CONTIGUOUS(means); - if (covars.has_value()) CHECK_CONTIGUOUS(covars.value()); - if (quats.has_value()) CHECK_CONTIGUOUS(quats.value()); - if (scales.has_value()) CHECK_CONTIGUOUS(scales.value()); + if (covars.has_value()) + CHECK_CONTIGUOUS(covars.value()); + if (quats.has_value()) + CHECK_CONTIGUOUS(quats.value()); + if (scales.has_value()) + CHECK_CONTIGUOUS(scales.value()); CHECK_CONTIGUOUS(viewmats); CHECK_CONTIGUOUS(Ks); CHECK_CONTIGUOUS(radii); CHECK_CONTIGUOUS(conics); - if (compensations.has_value()) CHECK_CONTIGUOUS(compensations.value()); + if (compensations.has_value()) + CHECK_CONTIGUOUS(compensations.value()); CHECK_CONTIGUOUS(v_means2d); CHECK_CONTIGUOUS(v_depths); CHECK_CONTIGUOUS(v_conics); - if (v_compensations.has_value()) CHECK_CONTIGUOUS(v_compensations.value()); + if (v_compensations.has_value()) + CHECK_CONTIGUOUS(v_compensations.value()); // Dimensions const uint32_t N = means.size(-2); @@ -53,28 +58,39 @@ projection_ewa_3dgs_fused_bwd( // 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()); + 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; + 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) { + 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, + 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, @@ -83,21 +99,32 @@ projection_ewa_3dgs_fused_bwd( camera_model, radii.data_ptr(), conics.data_ptr(), - compensations.has_value() ? compensations.value().data_ptr() : nullptr, + 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_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 + 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 ); - 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); diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp index 3f9c9f01..fd49d834 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp @@ -83,7 +83,9 @@ projection_ewa_3dgs_fused_fwd( sycl::range<1> global_range(num_work_groups * GSPLAT_N_THREADS); AT_DISPATCH_FLOATING_TYPES( - means.scalar_type(), "projection_ewa_3dgs_fused_fwd", [&] { + means.scalar_type(), + "projection_ewa_3dgs_fused_fwd", + [&] { auto e = d_queue.submit([&](sycl::handler &cgh) { FullyFusedProjectionFwdKernel kernel( B, diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp index 2650aaca..06d9e4f4 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp @@ -1,7 +1,7 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/PackedProjectionBwdKernel.hpp" namespace gsplat::xpu { @@ -31,13 +31,23 @@ projection_ewa_3dgs_packed_bwd( const at::Tensor v_conics, // [nnz, 3] const at::optional v_compensations, // [nnz] optional const bool viewmats_requires_grad, - const bool sparse_grad) { - - TORCH_CHECK(means.is_contiguous(), "Input 'means' tensor must be contiguous."); - TORCH_CHECK(viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous."); + const bool sparse_grad +) { + TORCH_CHECK( + means.is_contiguous(), "Input 'means' tensor must be contiguous." + ); + TORCH_CHECK( + viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous." + ); TORCH_CHECK(Ks.is_contiguous(), "Input 'Ks' tensor must be contiguous."); - TORCH_CHECK(batch_ids.is_contiguous(), "Input 'batch_ids' tensor must be contiguous."); - TORCH_CHECK(means.device().type() == at::kXPU, "Input tensors must be on XPU device."); + TORCH_CHECK( + batch_ids.is_contiguous(), + "Input 'batch_ids' tensor must be contiguous." + ); + TORCH_CHECK( + means.device().type() == at::kXPU, + "Input tensors must be on XPU device." + ); uint32_t N = means.size(-2); uint32_t C = viewmats.size(-3); @@ -70,43 +80,65 @@ projection_ewa_3dgs_packed_bwd( } if (nnz == 0) { - return std::make_tuple(v_means, v_covars, v_quats, v_scales, v_viewmats); + return std::make_tuple( + v_means, v_covars, v_quats, v_scales, v_viewmats + ); } - - auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + + 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::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(); - }); + 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); } diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp index 81402407..7a944ba5 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -1,7 +1,7 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/PackedProjectionFwdKernel.hpp" namespace gsplat::xpu { @@ -31,12 +31,19 @@ projection_ewa_3dgs_packed_fwd( const float far_plane, const float radius_clip, const bool calc_compensations, - const CameraModelType camera_model) { - - TORCH_CHECK(means.is_contiguous(), "Input 'means' tensor must be contiguous."); - TORCH_CHECK(viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous."); + const CameraModelType camera_model +) { + TORCH_CHECK( + means.is_contiguous(), "Input 'means' tensor must be contiguous." + ); + TORCH_CHECK( + viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous." + ); TORCH_CHECK(Ks.is_contiguous(), "Input 'Ks' tensor must be contiguous."); - TORCH_CHECK(means.device().type() == at::kXPU, "Input tensors must be on XPU device."); + TORCH_CHECK( + means.device().type() == at::kXPU, + "Input tensors must be on XPU device." + ); uint32_t N = means.size(-2); uint32_t C = viewmats.size(-3); @@ -64,60 +71,109 @@ projection_ewa_3dgs_packed_fwd( if (B == 0 || C == 0 || N == 0) { return std::make_tuple( - batch_ids, camera_ids, gaussian_ids, radii, means2d, depths, conics, indptr, compensations); + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + indptr, + compensations + ); } - + // --- Start of Correction --- // Allocate block_cnts as kInt, which the kernel expects. at::Tensor block_cnts = at::empty({(long)n_blocks}, int_opts); // --- End of Correction --- - auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + 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 - // --- Start of Correction --- - // Pass the int32_t pointer directly, no cast needed. - block_cnts.data_ptr(), - // --- End of Correction --- - nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr - )).wait(); - }); + 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 + // --- Start of Correction --- + // Pass the int32_t pointer directly, no cast needed. + block_cnts.data_ptr(), + // --- End of Correction --- + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr, + nullptr + ) + ) + .wait(); + } + ); // --- Start of Correction --- - // Perform inclusive scan on a kLong version of block_cnts to prevent overflow. + // 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); // --- End of Correction --- - + int64_t nnz = 0; if (n_blocks > 0) { - nnz = block_accum_inclusive.index({-1}).item(); + 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); + return std::make_tuple( + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + indptr, + compensations + ); } - - at::Tensor block_accum_exclusive = at::cat({at::zeros({1}, long_opts), block_accum_inclusive.slice(0, 0, n_blocks - 1)}); - at::Tensor block_accum_exclusive_int = block_accum_exclusive.to(at::kInt); + at::Tensor block_accum_exclusive = at::cat( + {at::zeros({1}, long_opts), + block_accum_inclusive.slice(0, 0, n_blocks - 1)} + ); + at::Tensor block_accum_exclusive_int = block_accum_exclusive.to(at::kInt); // Allocate final output tensors batch_ids = at::empty({nnz}, long_opts); @@ -130,34 +186,55 @@ projection_ewa_3dgs_packed_fwd( 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_exclusive_int.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(); - }); + 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_exclusive_int.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) { @@ -165,7 +242,16 @@ projection_ewa_3dgs_packed_fwd( } return std::make_tuple( - batch_ids, camera_ids, gaussian_ids, radii, means2d, depths, conics, indptr, compensations); + batch_ids, + camera_ids, + gaussian_ids, + radii, + means2d, + depths, + conics, + indptr, + 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 index 8b1deead..a19bc7fc 100644 --- a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp @@ -1,29 +1,27 @@ #include -#include "Ops.h" #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 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] + const at::Tensor v_means2d, // [..., C, N, 2] + const at::Tensor v_covars2d // [..., C, N, 2, 2] ) { - CHECK_CONTIGUOUS(means); CHECK_CONTIGUOUS(covars); CHECK_CONTIGUOUS(Ks); CHECK_CONTIGUOUS(v_means2d); CHECK_CONTIGUOUS(v_covars2d); - const uint32_t C = means.size(-3); const uint32_t N = means.size(-2); const uint32_t total_gaussians = means.numel() / 3; @@ -31,39 +29,36 @@ std::tuple projection_ewa_simple_bwd( 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; - + 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); - } - ); + 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); } diff --git a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp index 00441691..2a692816 100644 --- a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp @@ -1,11 +1,11 @@ #include -#include "Ops.h" #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] @@ -17,9 +17,16 @@ std::tuple projection_ewa_simple_fwd( CHECK_CONTIGUOUS(means); CHECK_CONTIGUOUS(covars); CHECK_CONTIGUOUS(Ks); - 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]"); + 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); @@ -37,35 +44,33 @@ std::tuple projection_ewa_simple_fwd( 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; - + 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); - } - ); + 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); } diff --git a/gsplat/sycl/src/projection_ut_3dgs_fused.cpp b/gsplat/sycl/src/projection_ut_3dgs_fused.cpp index 254a6f42..d4cbfd87 100644 --- a/gsplat/sycl/src/projection_ut_3dgs_fused.cpp +++ b/gsplat/sycl/src/projection_ut_3dgs_fused.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple< at::Tensor, at::Tensor, @@ -19,8 +19,8 @@ projection_ut_3dgs_fused( 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] + 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, @@ -32,10 +32,12 @@ projection_ut_3dgs_fused( // uncented transform const UnscentedTransformParameters ut_params, ShutterType rs_type, - const at::optional radial_coeffs, // [..., C, 6] or [..., C, 4] optional + 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 + 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"); } diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp index 1868cd15..2961d652 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -4,10 +4,10 @@ #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 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] @@ -20,7 +20,10 @@ std::tuple quat_scale_to_covar_preci_bwd( if (v_precis.has_value()) { CHECK_CONTIGUOUS(v_precis.value()); } - TORCH_CHECK(v_covars.has_value() || v_precis.has_value(), "Must provide gradients for at least one of covars or precis"); + 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); @@ -29,30 +32,27 @@ std::tuple quat_scale_to_covar_preci_bwd( if (N == 0) { return std::make_tuple(v_quats, v_scales); } - - auto& d_queue = at::xpu::getCurrentXPUStream().queue(); - + + 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); - 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); - } - ); + 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); + }); return std::make_tuple(v_quats, v_scales); } diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp index 8ab9ce49..b79f10ca 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -4,17 +4,20 @@ #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 at::Tensor quats, // [..., 4] + const at::Tensor scales, // [..., 3] const bool compute_covar, const bool compute_preci, const bool triu ) { CHECK_CONTIGUOUS(quats); CHECK_CONTIGUOUS(scales); - TORCH_CHECK(compute_covar || compute_preci, "Must compute at least one of covar or preci"); + 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(); @@ -47,28 +50,25 @@ std::tuple quat_scale_to_covar_preci_fwd( return std::make_tuple(covars, precis); } - auto& d_queue = at::xpu::getCurrentXPUStream().queue(); - + 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); - 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); - } - ); - + 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); + }); + return std::make_tuple(covars, precis); } diff --git a/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp b/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp index e1c6f224..71871778 100644 --- a/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp +++ b/gsplat/sycl/src/rasterize_to_indices_2dgs.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple rasterize_to_indices_2dgs( const uint32_t range_start, const uint32_t range_end, // iteration steps diff --git a/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp b/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp index a9594011..96726f30 100644 --- a/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp +++ b/gsplat/sycl/src/rasterize_to_indices_3dgs.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple rasterize_to_indices_3dgs( const uint32_t range_start, const uint32_t range_end, // iteration steps diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp index 85418bf7..2bc20ce0 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -1,7 +1,7 @@ #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" #include "kernels/RasterizeToPixels2DGSBwdKernel.hpp" namespace gsplat::xpu { @@ -11,50 +11,50 @@ 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, + 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, + 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, + 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, + 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 + 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(); + 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 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; @@ -66,57 +66,81 @@ void launch_rasterize_2dgs_bwd_kernel( 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); - } - ); + 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(); } @@ -132,32 +156,34 @@ std::tuple< 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] + 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] + 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] + 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] + 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 ) { // Check input tensors are contiguous @@ -178,11 +204,13 @@ rasterize_to_pixels_2dgs_bwd( CHECK_CONTIGUOUS(v_render_normals); CHECK_CONTIGUOUS(v_render_distort); CHECK_CONTIGUOUS(v_render_median); - if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); - if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); - + if (backgrounds.has_value()) + CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) + CHECK_CONTIGUOUS(masks.value()); + uint32_t channels = colors.size(-1); - + // Create output tensors auto options = means2d.options().dtype(torch::kFloat32); at::Tensor v_means2d_abs; @@ -195,20 +223,41 @@ rasterize_to_pixels_2dgs_bwd( 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 \ - ); \ +#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) { @@ -231,11 +280,11 @@ rasterize_to_pixels_2dgs_bwd( __GS__CALL_(257); __GS__CALL_(512); __GS__CALL_(513); - default: - TORCH_CHECK(false, "Unsupported number of channels: ", channels); + default: + TORCH_CHECK(false, "Unsupported number of channels: ", channels); } #undef __GS__CALL_ - + return std::make_tuple( v_means2d_abs, v_means2d, diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp index 35a6350d..3fc64713 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp @@ -5,30 +5,30 @@ namespace gsplat::xpu { -namespace { +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, + 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, + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, // forward outputs - const at::Tensor& render_alphas, - const at::Tensor& last_ids, + 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, + const at::Tensor &v_render_colors, + const at::Tensor &v_render_alphas, // options and derived params bool absgrad, bool packed, @@ -38,63 +38,88 @@ void launch_rasterize_bwd_kernel( 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 + 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(); - + 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::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); + + 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(); } @@ -103,12 +128,14 @@ void launch_rasterize_bwd_kernel( 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 + 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, @@ -117,11 +144,12 @@ rasterize_to_pixels_3dgs_bwd( 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] + 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] + 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 ) { @@ -135,8 +163,10 @@ rasterize_to_pixels_3dgs_bwd( CHECK_CONTIGUOUS(last_ids); CHECK_CONTIGUOUS(v_render_colors); CHECK_CONTIGUOUS(v_render_alphas); - if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); - if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + if (backgrounds.has_value()) + CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) + CHECK_CONTIGUOUS(masks.value()); TORCH_CHECK(means2d.dim() >= 2, "means2d must have at least 2 dimensions"); TORCH_CHECK(colors.dim() >= 2, "colors must have at least 2 dimensions"); @@ -149,22 +179,45 @@ rasterize_to_pixels_3dgs_bwd( 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 \ - ); \ + 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) { @@ -187,12 +240,14 @@ rasterize_to_pixels_3dgs_bwd( __GS_BWD_CALL_(257); __GS_BWD_CALL_(512); __GS_BWD_CALL_(513); - default: - TORCH_CHECK(false, "Unsupported number of channels: ", COLOR_DIM); + 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); + 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 index 32a1b1a5..2875f900 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp @@ -5,24 +5,24 @@ namespace gsplat::xpu { -namespace { +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, + 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, + const at::Tensor &tile_offsets, + const at::Tensor &flatten_ids, // other params bool packed, uint32_t C, @@ -30,59 +30,83 @@ void launch_rasterize_kernel( uint32_t tile_height, uint32_t tile_width, // outputs - at::Tensor& renders, - at::Tensor& alphas, - at::Tensor& last_ids + at::Tensor &renders, + at::Tensor &alphas, + at::Tensor &last_ids ) { - auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + 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::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 + 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 ); - cgh.parallel_for(range, kernel); } - ); + + 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 + 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, @@ -97,8 +121,10 @@ std::tuple rasterize_to_pixels_3dgs_fwd( CHECK_CONTIGUOUS(opacities); CHECK_CONTIGUOUS(tile_offsets); CHECK_CONTIGUOUS(flatten_ids); - if (backgrounds.has_value()) CHECK_CONTIGUOUS(backgrounds.value()); - if (masks.has_value()) CHECK_CONTIGUOUS(masks.value()); + if (backgrounds.has_value()) + CHECK_CONTIGUOUS(backgrounds.value()); + if (masks.has_value()) + CHECK_CONTIGUOUS(masks.value()); TORCH_CHECK(means2d.dim() >= 2, "means2d must have at least 2 dimensions"); TORCH_CHECK(colors.dim() >= 2, "colors must have at least 2 dimensions"); @@ -109,10 +135,12 @@ std::tuple rasterize_to_pixels_3dgs_fwd( 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 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}); @@ -127,14 +155,29 @@ std::tuple rasterize_to_pixels_3dgs_fwd( 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 \ - ); \ +#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) { @@ -157,8 +200,8 @@ std::tuple rasterize_to_pixels_3dgs_fwd( __GS__CALL_(257); __GS__CALL_(512); __GS__CALL_(513); - default: - TORCH_CHECK(false, "Unsupported number of channels: ", channels); + default: + TORCH_CHECK(false, "Unsupported number of channels: ", channels); } #undef __GS__CALL_ 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 index bc3faba3..e9233c5c 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_bwd.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple rasterize_to_pixels_from_world_3dgs_bwd( // Gaussian parameters @@ -15,30 +15,32 @@ rasterize_to_pixels_from_world_3dgs_bwd( 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] + 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::Tensor viewmats0, // [..., C, 4, 4] const at::optional - viewmats1, // [..., C, 4, 4] optional for rolling shutter - const at::Tensor Ks, // [..., C, 3, 3] + 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 + 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 + 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] + 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] + 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] 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 index 4c3d0758..e0bb0d64 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_from_world_3dgs_fwd.cpp @@ -1,11 +1,11 @@ - + #include -#include "Ops.h" #include "Common.h" +#include "Ops.h" + +namespace gsplat::xpu { -namespace gsplat::xpu { - std::tuple rasterize_to_pixels_from_world_3dgs_fwd( // Gaussian parameters @@ -15,24 +15,26 @@ rasterize_to_pixels_from_world_3dgs_fwd( 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] + 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::Tensor viewmats0, // [..., C, 4, 4] const at::optional - viewmats1, // [..., C, 4, 4] optional for rolling shutter - const at::Tensor Ks, // [..., C, 3, 3] + 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 + 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 + 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] diff --git a/gsplat/sycl/src/relocation.cpp b/gsplat/sycl/src/relocation.cpp index 079d6655..fcb31d8d 100644 --- a/gsplat/sycl/src/relocation.cpp +++ b/gsplat/sycl/src/relocation.cpp @@ -24,21 +24,21 @@ std::tuple relocation( at::Tensor new_scales = at::empty_like(scales); AT_DISPATCH_FLOATING_TYPES(opacities.scalar_type(), "relocation", ([&] { - auto &q = - c10::xpu::getCurrentXPUStream().queue(); - q.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() - ) - ) - .wait(); + 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); diff --git a/gsplat/sycl/src/spherical_harmonics_bwd.cpp b/gsplat/sycl/src/spherical_harmonics_bwd.cpp index 7845c74a..c2e4420d 100644 --- a/gsplat/sycl/src/spherical_harmonics_bwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_bwd.cpp @@ -8,10 +8,10 @@ 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::Tensor dirs, // [..., 3] + const at::Tensor coeffs, // [..., K, 3] const at::optional masks, // [...] - const at::Tensor v_colors, // [..., 3] + const at::Tensor v_colors, // [..., 3] bool compute_v_dirs ) { CHECK_CONTIGUOUS(dirs); @@ -28,36 +28,34 @@ std::tuple spherical_harmonics_bwd( 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()); + 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(); - + + 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); - 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); - } - ); + 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); + }); return std::make_tuple(v_coeffs, v_dirs); } diff --git a/gsplat/sycl/src/spherical_harmonics_fwd.cpp b/gsplat/sycl/src/spherical_harmonics_fwd.cpp index 86385de2..4393c716 100644 --- a/gsplat/sycl/src/spherical_harmonics_fwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_fwd.cpp @@ -3,7 +3,7 @@ #include "Ops.h" #include "kernels/ComputeShFwdKernel.hpp" -namespace gsplat::xpu { +namespace gsplat::xpu { at::Tensor spherical_harmonics_fwd( const uint32_t degrees_to_use, @@ -11,14 +11,27 @@ at::Tensor spherical_harmonics_fwd( const at::Tensor coeffs, // [..., K, 3] const at::optional masks // [...] ) { - TORCH_CHECK(dirs.is_contiguous(), "Input 'dirs' tensor must be contiguous."); - TORCH_CHECK(coeffs.is_contiguous(), "Input 'coeffs' tensor must be contiguous."); + TORCH_CHECK( + dirs.is_contiguous(), "Input 'dirs' tensor must be contiguous." + ); + TORCH_CHECK( + coeffs.is_contiguous(), "Input 'coeffs' tensor must be contiguous." + ); if (masks.has_value()) { - TORCH_CHECK(masks.value().is_contiguous(), "Input 'masks' tensor must be contiguous."); + TORCH_CHECK( + masks.value().is_contiguous(), + "Input 'masks' tensor must be contiguous." + ); } - - 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."); + + 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; @@ -29,28 +42,25 @@ at::Tensor spherical_harmonics_fwd( return colors; } - auto& d_queue = at::xpu::getCurrentXPUStream().queue(); + 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); - } - ); + + 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; diff --git a/setup.py b/setup.py index 987bc1ee..5d6e69cb 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,7 @@ try: import torch - has_xpu = has_cuda and hasattr(torch, 'xpu') and torch.xpu.is_available() + has_xpu = has_cuda and hasattr(torch, "xpu") and torch.xpu.is_available() except (ImportError, AttributeError): pass diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 241e36e1..8721888d 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -322,13 +322,7 @@ def test_rasterize_to_pixels_2dgs( normals.requires_grad = True densify.requires_grad = True - ( - render_colors, - render_alphas, - render_normals, - _, - _, - ) = rasterize_to_pixels_2dgs( + (render_colors, render_alphas, render_normals, _, _,) = rasterize_to_pixels_2dgs( means2d, ray_transforms, colors, From ed7352a36225f4735d154ff61501c4b20f3ee1c7 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Mon, 17 Nov 2025 13:54:19 -0800 Subject: [PATCH 32/56] Fix test_rasterization for packed, fix has_xpu check from previous commit. --- examples/simple_viewer.py | 2 +- .../include/kernels/PackedProjectionFwdKernel.hpp | 11 ++++++----- gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp | 8 +------- setup.py | 2 +- 4 files changed, 9 insertions(+), 14 deletions(-) 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/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp index 66b51867..26c22bed 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -28,7 +28,7 @@ template struct PackedProjectionFwdKernel { const T m_far_plane; const T m_radius_clip; const CameraModelType m_camera_model; - const int32_t *m_block_accum; // Packing helper for the second pass + const int64_t *m_block_accum; // Packing helper for the second pass // Outputs int32_t *m_block_cnts; @@ -60,7 +60,7 @@ template struct PackedProjectionFwdKernel { T far_plane, T radius_clip, CameraModelType camera_model, - const int32_t *block_accum, + const int64_t *block_accum, // outputs int32_t *block_cnts, int32_t *indptr, @@ -257,10 +257,10 @@ template struct PackedProjectionFwdKernel { } // --- Pass-specific logic --- - int32_t thread_data = static_cast(valid); 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. @@ -277,16 +277,17 @@ template struct PackedProjectionFwdKernel { } 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. - int32_t local_offset = sycl::exclusive_scan_over_group( + int64_t local_offset = sycl::exclusive_scan_over_group( group, thread_data, sycl::plus<>() ); if (valid) { - int32_t global_offset = local_offset; + int64_t global_offset = local_offset; if (block_idx > 0) { global_offset += m_block_accum[block_idx - 1]; } diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp index 7a944ba5..339c516f 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -169,12 +169,6 @@ projection_ewa_3dgs_packed_fwd( ); } - at::Tensor block_accum_exclusive = at::cat( - {at::zeros({1}, long_opts), - block_accum_inclusive.slice(0, 0, n_blocks - 1)} - ); - at::Tensor block_accum_exclusive_int = block_accum_exclusive.to(at::kInt); - // Allocate final output tensors batch_ids = at::empty({nnz}, long_opts); camera_ids = at::empty({nnz}, long_opts); @@ -218,7 +212,7 @@ projection_ewa_3dgs_packed_fwd( (scalar_t)far_plane, (scalar_t)radius_clip, camera_model, - block_accum_exclusive_int.data_ptr(), + block_accum_inclusive.data_ptr(), nullptr, // block_cnts indptr.data_ptr(), batch_ids.data_ptr(), diff --git a/setup.py b/setup.py index 5d6e69cb..7aab77a1 100644 --- a/setup.py +++ b/setup.py @@ -25,7 +25,7 @@ try: import torch - has_xpu = has_cuda and hasattr(torch, "xpu") and torch.xpu.is_available() + has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() except (ImportError, AttributeError): pass From b5bc4f082690abe17b151c0f7b4e5fc40ddd52f0 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Mon, 17 Nov 2025 14:12:36 -0800 Subject: [PATCH 33/56] #pragma once to follow gsplat convention --- gsplat/sycl/ext.cpp | 136 ++++++++++++------ .../{gsplat_sycl_utils.hpp => Sycl_utils.hpp} | 7 +- gsplat/sycl/include/helpers.hpp | 16 --- .../include/kernels/ComputeShBwdKernel.hpp | 7 +- .../include/kernels/ComputeShFwdKernel.hpp | 5 +- .../kernels/FullyFusedProjectionBwdKernel.hpp | 6 +- .../kernels/FullyFusedProjectionFwdKernel.hpp | 4 +- .../kernels/IsectOffsetEncodeKernel.hpp | 5 +- .../sycl/include/kernels/IsectTilesKernel.hpp | 5 +- .../kernels/PackedProjectionBwdKernel.hpp | 75 ++-------- .../kernels/PackedProjectionFwdKernel.hpp | 7 +- gsplat/sycl/include/kernels/ProjBwdKernel.hpp | 5 +- gsplat/sycl/include/kernels/ProjFwdKernel.hpp | 4 +- .../kernels/Projection2DGSFusedBwdKernel.hpp | 6 +- .../kernels/Projection2DGSFusedFwdKernel.hpp | 5 +- .../QuatScaleToCovarPreciBwdKernel.hpp | 5 +- .../QuatScaleToCovarPreciFwdKernel.hpp | 4 +- .../RasterizeToPixels2DGSBwdKernel.hpp | 7 +- .../RasterizeToPixels2DGSFwdKernel.hpp | 7 +- .../kernels/RasterizeToPixelsBwdKernel.hpp | 7 +- .../kernels/RasterizeToPixelsFwdKernel.hpp | 7 +- .../include/kernels/WorldToCamBwdKernel.hpp | 6 +- .../include/kernels/WorldToCamFwdKernel.hpp | 5 +- gsplat/sycl/include/proj.hpp | 5 +- gsplat/sycl/include/quat.hpp | 7 +- .../include/quat_scale_to_covar_preci.hpp | 5 +- gsplat/sycl/include/spherical_harmonics.hpp | 7 +- gsplat/sycl/include/transform.hpp | 7 +- gsplat/sycl/include/types.hpp | 7 +- gsplat/sycl/include/utils.hpp | 17 +-- .../src/quat_scale_to_covar_preci_bwd.cpp | 3 +- .../src/quat_scale_to_covar_preci_fwd.cpp | 3 +- gsplat/sycl/src/spherical_harmonics_bwd.cpp | 3 +- gsplat/sycl/src/spherical_harmonics_fwd.cpp | 2 +- 34 files changed, 153 insertions(+), 254 deletions(-) rename gsplat/sycl/include/{gsplat_sycl_utils.hpp => Sycl_utils.hpp} (96%) delete mode 100644 gsplat/sycl/include/helpers.hpp diff --git a/gsplat/sycl/ext.cpp b/gsplat/sycl/ext.cpp index 144a889f..74346ca0 100644 --- a/gsplat/sycl/ext.cpp +++ b/gsplat/sycl/ext.cpp @@ -1,76 +1,94 @@ #include -#include "Ops.h" #include "Cameras.h" +#include "Ops.h" PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - py::enum_< gsplat::xpu::CameraModelType>(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) + 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("null", &gsplat::xpu::null); m.def( - "quat_scale_to_covar_preci_fwd", & gsplat::xpu::quat_scale_to_covar_preci_fwd + "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 + "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("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("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("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_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 + "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 + "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 + &gsplat::xpu::projection_ewa_3dgs_packed_fwd ); m.def( "projection_ewa_3dgs_packed_bwd", - & gsplat::xpu::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 + "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 + "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("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("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 + "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 + "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("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); + 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") @@ -86,19 +104,47 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .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); + .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) + 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") + 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); + .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/gsplat_sycl_utils.hpp b/gsplat/sycl/include/Sycl_utils.hpp similarity index 96% rename from gsplat/sycl/include/gsplat_sycl_utils.hpp rename to gsplat/sycl/include/Sycl_utils.hpp index 2212587f..d312e797 100644 --- a/gsplat/sycl/include/gsplat_sycl_utils.hpp +++ b/gsplat/sycl/include/Sycl_utils.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SYCL_UTILS -#define GSPLAT_SYCL_UTILS +#pragma once #include @@ -68,6 +67,4 @@ template void gpuAtomicAddLocal(T &ref, const T &value) { sycl::access::address_space::local_space> protected_ref(ref); protected_ref.fetch_add(value); -} - -#endif \ No newline at end of file +} \ No newline at end of file diff --git a/gsplat/sycl/include/helpers.hpp b/gsplat/sycl/include/helpers.hpp deleted file mode 100644 index 5abd1c65..00000000 --- a/gsplat/sycl/include/helpers.hpp +++ /dev/null @@ -1,16 +0,0 @@ -#ifndef GSPLAT_SYCL_HELPERS_HPP -#define GSPLAT_SYCL_HELPERS_HPP - -#include - -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); -} - -#endif // GSPLAT_SYCL_HELPERS_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp index 6b958954..32f9650e 100644 --- a/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ComputeShBwdKernel.hpp @@ -1,9 +1,8 @@ -#ifndef ComputeShBwdKernel_HPP -#define ComputeShBwdKernel_HPP +#pragma once +#include "Sycl_utils.hpp" #include "spherical_harmonics.hpp" #include "types.hpp" -#include "utils.hpp" namespace gsplat::xpu { @@ -62,6 +61,4 @@ template struct ComputeShBwdKernel { } }; -#endif // ComputeShBwdKernel_HPP - } // 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 index 93def949..1abd3402 100644 --- a/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ComputeShFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef ComputeShFwdKernel_HPP -#define ComputeShFwdKernel_HPP +#pragma once #include "spherical_harmonics.hpp" @@ -46,6 +45,4 @@ template struct ComputeShFwdKernel { } }; -#endif // ComputeShFwdKernel_HPP - } // 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 index de5b29a8..4789bc25 100644 --- a/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionBwdKernel.hpp @@ -1,6 +1,6 @@ -#ifndef FullyFusedProjectionBwdKernel_HPP -#define FullyFusedProjectionBwdKernel_HPP +#pragma once +#include "Sycl_utils.hpp" #include "proj.hpp" #include "quat.hpp" #include "quat_scale_to_covar_preci.hpp" @@ -286,6 +286,4 @@ template struct FullyFusedProjectionBwdKernel { } }; -#endif // FullyFusedProjectionBwdKernel_HPP - } // 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 index 5bae3b52..951d3b5e 100644 --- a/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/FullyFusedProjectionFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef FullyFusedProjectionFwdKernel_HPP -#define FullyFusedProjectionFwdKernel_HPP +#pragma once #include "proj.hpp" #include "quat_scale_to_covar_preci.hpp" @@ -252,6 +251,5 @@ template struct FullyFusedProjectionFwdKernel { } } }; -#endif // FullyFusedProjectionFwdKernel_HPP } // 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 index b64f4f62..92f968b0 100644 --- a/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp +++ b/gsplat/sycl/include/kernels/IsectOffsetEncodeKernel.hpp @@ -1,5 +1,4 @@ -#ifndef IsectOffsetEncodeKernel_HPP -#define IsectOffsetEncodeKernel_HPP +#pragma once namespace gsplat::xpu { @@ -63,6 +62,4 @@ struct IsectOffsetEncodeKernel { } }; -#endif - } // 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 index ee91751f..d5cb93a7 100644 --- a/gsplat/sycl/include/kernels/IsectTilesKernel.hpp +++ b/gsplat/sycl/include/kernels/IsectTilesKernel.hpp @@ -1,5 +1,4 @@ -#ifndef IsectTilesKernel_HPP -#define IsectTilesKernel_HPP +#pragma once #include "transform.hpp" #include "types.hpp" @@ -149,6 +148,4 @@ template struct IsectTilesKernel { } }; -#endif // IsectTilesKernel_HPP - } // 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 index 53087853..0b8b96bd 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionBwdKernel.hpp @@ -1,6 +1,6 @@ -#ifndef PackedProjectionBwdKernel_HPP -#define PackedProjectionBwdKernel_HPP +#pragma once +#include "Sycl_utils.hpp" #include "proj.hpp" #include "quat.hpp" #include "quat_scale_to_covar_preci.hpp" @@ -273,46 +273,17 @@ template struct PackedProjectionBwdKernel { 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) { - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device> - ref(v_means_out[i]); - ref.fetch_add(v_mean[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; - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_covars_out[0]) - .fetch_add(v_covar[0][0]); - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_covars_out[1]) - .fetch_add(v_covar[0][1] + v_covar[1][0]); - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_covars_out[2]) - .fetch_add(v_covar[0][2] + v_covar[2][0]); - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_covars_out[3]) - .fetch_add(v_covar[1][1]); - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_covars_out[4]) - .fetch_add(v_covar[1][2] + v_covar[2][1]); - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_covars_out[5]) - .fetch_add(v_covar[2][2]); + 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); @@ -323,17 +294,9 @@ template struct PackedProjectionBwdKernel { 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) - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_quats_out[i]) - .fetch_add(v_quat[i]); + gpuAtomicAdd(&v_quats_out[i], v_quat[i]); for (int i = 0; i < 3; ++i) - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device>(v_scales_out[i]) - .fetch_add(v_scale[i]); + gpuAtomicAdd(&v_scales_out[i], v_scale[i]); } } @@ -342,24 +305,12 @@ template struct PackedProjectionBwdKernel { 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 - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device> - ref(v_viewmats_out[i * 4 + j]); - ref.fetch_add(v_R[j][i]); + gpuAtomicAdd(&v_viewmats_out[i * 4 + j], v_R[j][i]); } - sycl::atomic_ref< - T, - sycl::memory_order::relaxed, - sycl::memory_scope::device> - ref(v_viewmats_out[i * 4 + 3]); - ref.fetch_add(v_t[i]); + gpuAtomicAdd(&v_viewmats_out[i * 4 + 3], v_t[i]); } } } }; } // namespace gsplat::xpu - -#endif // PackedProjectionBwdKernel_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp index 26c22bed..8ce95ca5 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef PackedProjectionFwdKernel_HPP -#define PackedProjectionFwdKernel_HPP +#pragma once #include "proj.hpp" #include "quat_scale_to_covar_preci.hpp" @@ -324,6 +323,4 @@ template struct PackedProjectionFwdKernel { } }; -} // namespace gsplat::xpu - -#endif // PackedProjectionFwdKernel_HPP \ No newline at end of file +} // 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 index b75a6eb7..5198f85f 100644 --- a/gsplat/sycl/include/kernels/ProjBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ProjBwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef ProjBwdKernel_HPP -#define ProjBwdKernel_HPP +#pragma once #include "Common.h" #include "proj.hpp" @@ -135,6 +134,4 @@ template struct ProjBwdKernel { } }; -#endif // ProjBwdKernel_HPP - } // 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 index d7b28ae4..411e1c01 100644 --- a/gsplat/sycl/include/kernels/ProjFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/ProjFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef ProjFwdKernel_HPP -#define ProjFwdKernel_HPP +#pragma once #include "Common.h" #include "proj.hpp" @@ -90,6 +89,5 @@ template struct ProjFwdKernel { } } }; -#endif // ProjFwdKernel_HPP } // 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 index 4a05409f..d938a035 100644 --- a/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedBwdKernel.hpp @@ -1,6 +1,6 @@ -#ifndef Projection2DGSFusedBwdKernel_HPP -#define Projection2DGSFusedBwdKernel_HPP +#pragma once +#include "Sycl_utils.hpp" #include "quat_scale_to_covar_preci.hpp" #include "transform.hpp" #include "utils.hpp" @@ -304,5 +304,3 @@ template struct Projection2DGSFusedBwdKernel { }; } // namespace gsplat::xpu - -#endif // Projection2DGSFusedBwdKernel_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp index cf9b28fb..fb97c166 100644 --- a/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/Projection2DGSFusedFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef Projection2DGSFusedFwdKernel_HPP -#define Projection2DGSFusedFwdKernel_HPP +#pragma once #include "quat_scale_to_covar_preci.hpp" #include "transform.hpp" @@ -219,5 +218,3 @@ template struct Projection2DGSFusedFwdKernel { }; } // namespace gsplat::xpu - -#endif // Projection2DGSFusedFwdKernel_HPP \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp index 667fc721..8c7e2b54 100644 --- a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciBwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef QuatScaleToCovarPreciBwdKernel_HPP -#define QuatScaleToCovarPreciBwdKernel_HPP +#pragma once #include "quat_scale_to_covar_preci.hpp" @@ -112,6 +111,4 @@ template struct QuatScaleToCovarPreciBwdKernel { } }; -#endif // QuatScaleToCovarPreciBwdKernel_HPP - } // 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 index 8e587b25..fef9cc63 100644 --- a/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/QuatScaleToCovarPreciFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef QuatScaleToCovarPreciFwdKernel_HPP -#define QuatScaleToCovarPreciFwdKernel_HPP +#pragma once #include "quat_scale_to_covar_preci.hpp" @@ -89,6 +88,5 @@ template struct QuatScaleToCovarPreciFwdKernel { } } }; -#endif // QuatScaleToCovarPreciFwdKernel_HPP } // 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 index fdf2dbab..f00108b6 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSBwdKernel.hpp @@ -1,7 +1,6 @@ -#ifndef RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP -#define RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP +#pragma once -#include "gsplat_sycl_utils.hpp" +#include "Sycl_utils.hpp" #include "types.hpp" #include @@ -714,5 +713,3 @@ template struct RasterizeToPixels2DGSBwdKernel { }; } // namespace gsplat::xpu - -#endif // RASTERIZE_TO_PIXELS_2DGS_BWD_KERNEL_HPP diff --git a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp index 06105089..bbf5201c 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixels2DGSFwdKernel.hpp @@ -1,7 +1,6 @@ -#ifndef RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP -#define RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP +#pragma once -#include "gsplat_sycl_utils.hpp" +#include "Sycl_utils.hpp" #include "types.hpp" #include @@ -411,5 +410,3 @@ template struct RasterizeToPixels2DGSFwdKernel { }; } // namespace gsplat::xpu - -#endif // RASTERIZE_TO_PIXELS_2DGS_FWD_KERNEL_HPP diff --git a/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp index f0ff4fbc..94fc968e 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsBwdKernel.hpp @@ -1,7 +1,6 @@ -#ifndef RasterizeToPixelsBwdKernel_HPP -#define RasterizeToPixelsBwdKernel_HPP +#pragma once -#include "gsplat_sycl_utils.hpp" +#include "Sycl_utils.hpp" #include "types.hpp" #include @@ -445,6 +444,4 @@ struct RasterizeToPixelsBwdKernel { } }; -#endif // RasterizeToPixelsBwdKernel_HPP - } // 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 index 8528c1fa..b3ae0a6a 100644 --- a/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/RasterizeToPixelsFwdKernel.hpp @@ -1,7 +1,6 @@ -#ifndef RasterizeToPixelsFwdKernel_HPP -#define RasterizeToPixelsFwdKernel_HPP +#pragma once -#include "gsplat_sycl_utils.hpp" +#include "Sycl_utils.hpp" #include "types.hpp" namespace gsplat::xpu { @@ -320,6 +319,4 @@ struct RasterizeToPixelsFwdKernel { } }; -#endif // RasterizeToPixelsFwdKernel_HPP - } // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp index 0ef6613c..0d8c4584 100644 --- a/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp +++ b/gsplat/sycl/include/kernels/WorldToCamBwdKernel.hpp @@ -1,6 +1,6 @@ -#ifndef WorldToCamBwdKernel_HPP -#define WorldToCamBwdKernel_HPP +#pragma once +#include "Sycl_utils.hpp" #include "transform.hpp" #include "types.hpp" #include "utils.hpp" @@ -117,6 +117,4 @@ template struct WorldToCamBwdKernel { } }; -#endif // WorldToCamBwdKernel_HPP - } // 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 index 59829a48..d272f63d 100644 --- a/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/WorldToCamFwdKernel.hpp @@ -1,5 +1,4 @@ -#ifndef WorldToCamFwdKernel_HPP -#define WorldToCamFwdKernel_HPP +#pragma once /**************************************************************************** * World to Camera Transformation Forward Pass @@ -90,6 +89,4 @@ template struct WorldToCamFwdKernel { } }; -#endif // WorldToCamFwdKernel_HPP - } // namespace gsplat::xpu \ No newline at end of file diff --git a/gsplat/sycl/include/proj.hpp b/gsplat/sycl/include/proj.hpp index 60a6de9c..4ed9890f 100644 --- a/gsplat/sycl/include/proj.hpp +++ b/gsplat/sycl/include/proj.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SYCL_PROJ_HPP -#define GSPLAT_SYCL_PROJ_HPP +#pragma once #include "types.hpp" @@ -339,5 +338,3 @@ inline void fisheye_proj_vjp( v_mean3d.y += dL_dty_raw; v_mean3d.z += dL_dtz_raw; } - -#endif // GSPLAT_SYCL_PROJ_HPP diff --git a/gsplat/sycl/include/quat.hpp b/gsplat/sycl/include/quat.hpp index 2c125814..9282bd81 100644 --- a/gsplat/sycl/include/quat.hpp +++ b/gsplat/sycl/include/quat.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SYCL_QUAT_HPP -#define GSPLAT_SYCL_QUAT_HPP +#pragma once #include "types.hpp" @@ -51,6 +50,4 @@ quat_to_rotmat_vjp(const vec4 quat, const mat3 v_R, vec4 &v_quat) { vec4 quat_n = vec4(w, x, y, z); v_quat += (v_quat_n - glm::dot(v_quat_n, quat_n) * quat_n) * inv_norm; -} - -#endif // GSPLAT_SYCL_QUAT_HPP +} \ 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 index f1f43a74..42ddb7bd 100644 --- a/gsplat/sycl/include/quat_scale_to_covar_preci.hpp +++ b/gsplat/sycl/include/quat_scale_to_covar_preci.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP -#define GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP +#pragma once #include "quat.hpp" #include "types.hpp" @@ -120,5 +119,3 @@ inline void quat_scale_to_preci_vjp( -sz * sz * (R[2][0] * v_M[2][0] + R[2][1] * v_M[2][1] + R[2][2] * v_M[2][2]); } - -#endif // GSPLAT_SYCL_QUAT_SCALE_TO_COVAR_PRECI_HPP diff --git a/gsplat/sycl/include/spherical_harmonics.hpp b/gsplat/sycl/include/spherical_harmonics.hpp index a94ca68b..9efe6207 100644 --- a/gsplat/sycl/include/spherical_harmonics.hpp +++ b/gsplat/sycl/include/spherical_harmonics.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP -#define GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP +#pragma once #include "types.hpp" @@ -355,6 +354,4 @@ inline void sh_coeffs_to_color_fast_vjp( v_dir->y = v_d.y; v_dir->z = v_d.z; } -} - -#endif // GSPLAT_SPHERICAL_HARMONICS_SYCL_HPP \ No newline at end of file +} \ No newline at end of file diff --git a/gsplat/sycl/include/transform.hpp b/gsplat/sycl/include/transform.hpp index e8e281eb..917c8a54 100644 --- a/gsplat/sycl/include/transform.hpp +++ b/gsplat/sycl/include/transform.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SYCL_TRANSFORM_HPP -#define GSPLAT_SYCL_TRANSFORM_HPP +#pragma once #include "types.hpp" @@ -64,6 +63,4 @@ inline void covar_world_to_cam_vjp( 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; -} - -#endif // GSPLAT_SYCL_TRANSFORM_HPP +} \ No newline at end of file diff --git a/gsplat/sycl/include/types.hpp b/gsplat/sycl/include/types.hpp index 70f18a4e..109b5536 100644 --- a/gsplat/sycl/include/types.hpp +++ b/gsplat/sycl/include/types.hpp @@ -1,5 +1,4 @@ -#ifndef GSPLAT_SYCL_TYPES_HPP -#define GSPLAT_SYCL_TYPES_HPP +#pragma once #include @@ -15,6 +14,4 @@ template using mat3 = glm::mat<3, 3, T>; template using mat4 = glm::mat<4, 4, T>; -template using mat3x2 = glm::mat<3, 2, T>; - -#endif // GSPLAT_SYCL_TYPES_HPP \ No newline at end of file +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 index 308537a0..0bdd1059 100644 --- a/gsplat/sycl/include/utils.hpp +++ b/gsplat/sycl/include/utils.hpp @@ -1,20 +1,9 @@ -#ifndef GSPLAT_SYCL_UTILS_HPP -#define GSPLAT_SYCL_UTILS_HPP +#pragma once #include "types.hpp" #include -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 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) { @@ -80,6 +69,4 @@ inline void add_blur_vjp( 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); -} - -#endif // GSPLAT_SYCL_UTILS_HPP +} \ 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 index 2961d652..9a3d1eba 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -40,7 +40,7 @@ std::tuple quat_scale_to_covar_preci_bwd( sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); sycl::nd_range<1> range(globalRange, localRange); - d_queue.submit([&](sycl::handler &cgh) { + auto e = d_queue.submit([&](sycl::handler &cgh) { QuatScaleToCovarPreciBwdKernel kernel( N, quats.data_ptr(), @@ -53,6 +53,7 @@ std::tuple quat_scale_to_covar_preci_bwd( ); cgh.parallel_for(range, kernel); }); + e.wait(); return std::make_tuple(v_quats, v_scales); } diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp index b79f10ca..d20be021 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -57,7 +57,7 @@ std::tuple quat_scale_to_covar_preci_fwd( sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); sycl::nd_range<1> range(globalRange, localRange); - d_queue.submit([&](sycl::handler &cgh) { + auto e = d_queue.submit([&](sycl::handler &cgh) { QuatScaleToCovarPreciFwdKernel kernel( N, quats.data_ptr(), @@ -68,6 +68,7 @@ std::tuple quat_scale_to_covar_preci_fwd( ); cgh.parallel_for(range, kernel); }); + e.wait(); return std::make_tuple(covars, precis); } diff --git a/gsplat/sycl/src/spherical_harmonics_bwd.cpp b/gsplat/sycl/src/spherical_harmonics_bwd.cpp index c2e4420d..ccba983d 100644 --- a/gsplat/sycl/src/spherical_harmonics_bwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_bwd.cpp @@ -42,7 +42,7 @@ std::tuple spherical_harmonics_bwd( sycl::range<1> globalRange(GSPLAT_N_THREADS * numWorkGrps); sycl::nd_range<1> range(globalRange, localRange); - d_queue.submit([&](sycl::handler &cgh) { + auto e = d_queue.submit([&](sycl::handler &cgh) { ComputeShBwdKernel kernel( N, K, @@ -56,6 +56,7 @@ std::tuple spherical_harmonics_bwd( ); cgh.parallel_for(range, kernel); }); + e.wait(); return std::make_tuple(v_coeffs, v_dirs); } diff --git a/gsplat/sycl/src/spherical_harmonics_fwd.cpp b/gsplat/sycl/src/spherical_harmonics_fwd.cpp index 4393c716..53f11cff 100644 --- a/gsplat/sycl/src/spherical_harmonics_fwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_fwd.cpp @@ -61,8 +61,8 @@ at::Tensor spherical_harmonics_fwd( ); cgh.parallel_for(range, kernel); }); - e.wait(); + return colors; } From 72fcd3a4b78e7c85e3ec322bedd8c888039e45cb Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Mon, 17 Nov 2025 22:23:13 -0800 Subject: [PATCH 34/56] Common _wrapper.py for cuda and sycl --- formatter.sh | 4 + gsplat/__init__.py | 55 +- gsplat/_torch_impl.py | 14 +- gsplat/_torch_impl_2dgs.py | 6 +- gsplat/{cuda => }/_wrapper.py | 104 +- gsplat/optimizers/selective_adam.py | 2 +- gsplat/relocation.py | 10 +- gsplat/rendering.py | 41 +- gsplat/strategy/default.py | 1 - gsplat/strategy/ops.py | 16 +- gsplat/sycl/_wrapper.py | 2612 --------------------------- setup.py | 2 +- tests/test_2dgs.py | 3 +- 13 files changed, 106 insertions(+), 2764 deletions(-) rename gsplat/{cuda => }/_wrapper.py (96%) delete mode 100644 gsplat/sycl/_wrapper.py 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 9df5004f..76f98316 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -6,29 +6,10 @@ torch_acc = torch.cpu _force_backend = os.getenv("GSPLAT_BACKEND", "").lower() -from .cuda._wrapper import ( # Default to CUDA imports, works even if no CUDA is available - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - proj, - quat_scale_to_covar_preci, - rasterize_to_indices_in_range, - rasterize_to_indices_in_range_2dgs, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - world_to_cam, -) - if _force_backend == "cuda" or (_force_backend == "" and torch.cuda.is_available()): BACKEND = "cuda" torch_acc = torch.cuda print("gsplat: Using CUDA backend.", file=sys.stderr) - # Functions already imported above if ( not BACKEND @@ -37,29 +18,27 @@ and hasattr(torch, "xpu") and torch.xpu.is_available() ): - from .sycl._wrapper import ( # Overwrite imports for SYCL backend - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - proj, - quat_scale_to_covar_preci, - rasterize_to_indices_in_range, - rasterize_to_indices_in_range_2dgs, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - world_to_cam, - ) - 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, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + proj, + quat_scale_to_covar_preci, + rasterize_to_indices_in_range, + rasterize_to_indices_in_range_2dgs, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, + world_to_cam, +) from .compression import PngCompression from .exporter import export_splats from .optimizers import SelectiveAdam diff --git a/gsplat/_torch_impl.py b/gsplat/_torch_impl.py index ab20ab5f..f06ca9df 100644 --- a/gsplat/_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 @@ -639,7 +639,7 @@ def _rasterize_to_pixels( This function requires the `nerfacc` package to be installed. Please install it using the following command `pip install nerfacc`. """ - from .cuda._wrapper import rasterize_to_indices_in_range + from ._wrapper import rasterize_to_indices_in_range image_dims = means2d.shape[:-2] channels = colors.shape[-1] @@ -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/_torch_impl_2dgs.py b/gsplat/_torch_impl_2dgs.py index 4f0fd4c3..7f3a0ab8 100644 --- a/gsplat/_torch_impl_2dgs.py +++ b/gsplat/_torch_impl_2dgs.py @@ -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 @@ -231,7 +231,7 @@ def _rasterize_to_pixels_2dgs( This function requires the `nerfacc` package to be installed. Please install it using the following command `pip install nerfacc`. """ - from .cuda._wrapper import rasterize_to_indices_in_range_2dgs + from ._wrapper import rasterize_to_indices_in_range_2dgs image_dims = means2d.shape[:-2] channels = colors.shape[-1] 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 50ba0f03..5d2f03c5 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 @@ -112,10 +118,10 @@ def world_to_cam( - **Gaussian means in camera coordinate system**. [..., C, N, 3] - **Gaussian covariances in camera coordinate system**. [..., C, N, 3, 3] """ - from .._torch_impl import _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 ) @@ -324,7 +330,7 @@ def fully_fused_projection( .. 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: @@ -501,7 +507,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(), @@ -535,7 +541,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 ) @@ -914,7 +920,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(), @@ -944,7 +950,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) @@ -963,7 +969,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, @@ -990,11 +996,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, @@ -1014,7 +1020,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, @@ -1053,12 +1059,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, @@ -1108,7 +1114,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, @@ -1214,9 +1220,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(), @@ -1269,7 +1275,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, @@ -1335,7 +1341,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, @@ -1413,7 +1419,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 = ( @@ -1422,7 +1428,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, @@ -1511,7 +1517,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, @@ -1605,7 +1611,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()}" ) @@ -1619,7 +1625,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 @@ -1701,7 +1707,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, @@ -1805,7 +1811,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) @@ -1819,7 +1825,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, @@ -1958,7 +1964,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, @@ -2004,7 +2010,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, @@ -2075,7 +2081,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, @@ -2140,7 +2146,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, @@ -2292,7 +2298,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], @@ -2420,7 +2426,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(), @@ -2468,7 +2474,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, @@ -2554,7 +2560,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, @@ -2579,7 +2585,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/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/relocation.py b/gsplat/relocation.py index 92b6ce93..0f6538b7 100644 --- a/gsplat/relocation.py +++ b/gsplat/relocation.py @@ -5,13 +5,7 @@ from torch import Tensor from . import BACKEND - -# Now, conditionally import the functions based on the detected backend. -if BACKEND == "cuda": - from .cuda._wrapper import _make_lazy_cuda_func as _make_lazy_func -elif BACKEND == "sycl": - from .sycl._wrapper import _make_lazy_sycl_func as _make_lazy_func - +from ._wrapper import _make_lazy_device_func def compute_relocation( opacities: Tensor, # [N] @@ -49,7 +43,7 @@ def compute_relocation( ratios.clamp_(min=1, max=n_max) ratios = ratios.int().contiguous() - new_opacities, new_scales = _make_lazy_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 8a839581..8dd06f5a 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -8,35 +8,18 @@ from typing_extensions import Literal from . import BACKEND - -# Now, conditionally import the functions based on the detected backend. -if BACKEND == "sycl": - from .sycl._wrapper import ( - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - ) -else: # CUDA or no backend (e.g., CPU only for docs and testing) - from .cuda._wrapper import ( - RollingShutterType, - fully_fused_projection, - fully_fused_projection_2dgs, - fully_fused_projection_with_ut, - isect_offset_encode, - isect_tiles, - rasterize_to_pixels, - rasterize_to_pixels_2dgs, - rasterize_to_pixels_eval3d, - spherical_harmonics, - ) - +from ._wrapper import ( + RollingShutterType, + fully_fused_projection, + fully_fused_projection_2dgs, + fully_fused_projection_with_ut, + isect_offset_encode, + isect_tiles, + rasterize_to_pixels, + rasterize_to_pixels_2dgs, + rasterize_to_pixels_eval3d, + spherical_harmonics, +) from .distributed import ( all_gather_int32, all_gather_tensor_list, diff --git a/gsplat/strategy/default.py b/gsplat/strategy/default.py index 8b3a6b18..2604f4c1 100644 --- a/gsplat/strategy/default.py +++ b/gsplat/strategy/default.py @@ -192,7 +192,6 @@ def step_post_backward( if self.refine_scale2d_stop_iter > 0: state["radii"].zero_() torch_acc.empty_cache() - print(f"Empty cache after step {step}", flush=True) if step % self.reset_every == 0 & step > 0: reset_opa( diff --git a/gsplat/strategy/ops.py b/gsplat/strategy/ops.py index 7a7aa66d..9f6892d8 100644 --- a/gsplat/strategy/ops.py +++ b/gsplat/strategy/ops.py @@ -5,19 +5,9 @@ import torch.nn.functional as F from torch import Tensor -from gsplat import BACKEND - -if BACKEND == "sycl": - from gsplat.sycl._wrapper import ( - quat_scale_to_covar_preci, - ) -else: # BACKEND == "cuda" or None - from gsplat.cuda._wrapper 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/_wrapper.py b/gsplat/sycl/_wrapper.py deleted file mode 100644 index 8e188e8b..00000000 --- a/gsplat/sycl/_wrapper.py +++ /dev/null @@ -1,2612 +0,0 @@ -import math -import warnings -from dataclasses import dataclass -from enum import Enum -from typing import Any, Callable, Optional, Tuple - -import torch -from torch import Tensor -from typing_extensions import Literal - - -def _make_lazy_sycl_func(name: str) -> Callable: - """Creates a lazy-loading function for the SYCL backend.""" - - def call_sycl(*args, **kwargs): - # pylint: disable=import-outside-toplevel - from ._backend import _C - - return getattr(_C, name)(*args, **kwargs) - - return call_sycl - - -def _make_lazy_sycl_obj(name: str) -> Any: - """Creates a lazy-loading object accessor for the SYCL backend.""" - # pylint: disable=import-outside-toplevel - from ._backend import _C - - obj = _C - for name_split in name.split("."): - obj = getattr(obj, name_split) - return obj - - -class RollingShutterType(Enum): - ROLLING_TOP_TO_BOTTOM = 0 - ROLLING_LEFT_TO_RIGHT = 1 - ROLLING_BOTTOM_TO_TOP = 2 - ROLLING_RIGHT_TO_LEFT = 3 - GLOBAL = 4 - - def to_cpp(self) -> Any: - return _make_lazy_sycl_obj(f"ShutterType.{self.name}") - - -@dataclass -class UnscentedTransformParameters: - # Sigma point parameters (see Gustafsson and Hendeby 2012, Wan and van der Merwe 2000) - alpha: float = 0.1 - beta: float = 2.0 - kappa: float = 0.0 - # Parameters controlling validity of the unscented transform results. Default 0.1 - # is 10% margin. - in_image_margin_factor: float = 0.1 - # True: all sigma points must be valid - require_all_sigma_points_valid: bool = True - - def to_cpp(self) -> Any: - p = _make_lazy_sycl_obj("UnscentedTransformParameters")() - p.alpha = self.alpha - p.beta = self.beta - p.kappa = self.kappa - p.in_image_margin_factor = self.in_image_margin_factor - p.require_all_sigma_points_valid = self.require_all_sigma_points_valid - return p - - -@dataclass -class FThetaPolynomialType(Enum): - PIXELDIST_TO_ANGLE = 0 - ANGLE_TO_PIXELDIST = 1 - - def to_cpp(self) -> Any: - return _make_lazy_sycl_obj(f"FThetaPolynomialType.{self.name}") - - -@dataclass -class FThetaCameraDistortionParameters: - reference_poly: FThetaPolynomialType - pixeldist_to_angle_poly: Tuple[float, float, float, float, float, float] # [6] - angle_to_pixeldist_poly: Tuple[float, float, float, float, float, float] # [6] - max_angle: float - linear_cde: Tuple[float, float, float] # [3] - - def to_cpp(self) -> Any: - p = _make_lazy_sycl_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 - p.max_angle = self.max_angle - p.linear_cde = self.linear_cde - return p - - @classmethod - def to_cpp_default(cls) -> Any: - p = _make_lazy_sycl_obj("FThetaCameraDistortionParameters")() - return p - - -def world_to_cam( - means: Tensor, # [..., N, 3] - covars: Tensor, # [..., N, 3, 3] - viewmats: Tensor, # [..., C, 4, 4] -) -> Tuple[Tensor, Tensor]: - """Transforms Gaussians from world to camera coordinate system. - - Args: - means: Gaussian means. [..., N, 3] - covars: Gaussian covariances. [..., N, 3, 3] - viewmats: World-to-camera transformation matrices. [..., C, 4, 4] - - Returns: - A tuple: - - - **Gaussian means in camera coordinate system**. [..., C, N, 3] - - **Gaussian covariances in camera coordinate system**. [..., C, N, 3, 3] - """ - from .._torch_impl import _world_to_cam - - warnings.warn( - "world_to_cam() is removed from the sycl 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, - ) - batch_dims = means.shape[:-2] - N = means.shape[-2] - C = viewmats.shape[-3] - assert means.shape == batch_dims + (N, 3), means.shape - assert covars.shape == batch_dims + (N, 3, 3), covars.shape - assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape - means = means.contiguous() - covars = covars.contiguous() - viewmats = viewmats.contiguous() - return _world_to_cam(means, covars, viewmats) - - -def adam( - param: Tensor, - param_grad: Tensor, - exp_avg: Tensor, - exp_avg_sq: Tensor, - valid: Tensor, - lr: float, - b1: float, - b2: float, - eps: float, -) -> None: - _make_lazy_sycl_func("adam")( - param, param_grad, exp_avg, exp_avg_sq, valid, lr, b1, b2, eps - ) - - -def spherical_harmonics( - degrees_to_use: int, - dirs: Tensor, # [..., 3] - coeffs: Tensor, # [..., K, 3] - masks: Optional[Tensor] = None, # [...,] -) -> Tensor: - """Computes spherical harmonics. - - Args: - degrees_to_use: The degree to be used. - dirs: Directions. [..., 3] - coeffs: Coefficients. [..., K, 3] - masks: Optional boolen masks to skip some computation. [...,] Default: None. - - Returns: - Spherical harmonics. [..., 3] - """ - assert (degrees_to_use + 1) ** 2 <= coeffs.shape[-2], coeffs.shape - batch_dims = dirs.shape[:-1] - assert dirs.shape == batch_dims + (3,), dirs.shape - assert ( - (len(coeffs.shape) == len(batch_dims) + 2) - and coeffs.shape[:-2] == batch_dims - and coeffs.shape[-1] == 3 - ), coeffs.shape - if masks is not None: - assert masks.shape == batch_dims, masks.shape - masks = masks.contiguous() - return _SphericalHarmonics.apply( - degrees_to_use, dirs.contiguous(), coeffs.contiguous(), masks - ) - - -def quat_scale_to_covar_preci( - quats: Tensor, # [..., 4], - scales: Tensor, # [..., 3], - compute_covar: bool = True, - compute_preci: bool = True, - triu: bool = False, -) -> Tuple[Optional[Tensor], Optional[Tensor]]: - """Converts quaternions and scales to covariance and precision matrices. - - Args: - quats: Quaternions (No need to be normalized). [..., 4] - scales: Scales. [..., 3] - compute_covar: Whether to compute covariance matrices. Default: True. If False, - the returned covariance matrices will be None. - compute_preci: Whether to compute precision matrices. Default: True. If False, - the returned precision matrices will be None. - triu: If True, the return matrices will be upper triangular. Default: False. - - Returns: - A tuple: - - - **Covariance matrices**. If `triu` is True the returned shape is [..., 6], otherwise [..., 3, 3]. - - **Precision matrices**. If `triu` is True the returned shape is [..., 6], otherwise [..., 3, 3]. - """ - batch_dims = quats.shape[:-1] - assert quats.shape == batch_dims + (4,), quats.shape - assert scales.shape == batch_dims + (3,), scales.shape - quats = quats.contiguous() - scales = scales.contiguous() - covars, precis = _QuatScaleToCovarPreci.apply( - quats, scales, compute_covar, compute_preci, triu - ) - return covars if compute_covar else None, precis if compute_preci else None - - -def persp_proj( - means: Tensor, # [..., C, N, 3] - covars: Tensor, # [..., C, N, 3, 3] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, -) -> Tuple[Tensor, Tensor]: - """Perspective projection on Gaussians. - DEPRECATED: please use `proj` with `ortho=False` instead. - - Args: - means: Gaussian means. [..., C, N, 3] - covars: Gaussian covariances. [..., C, N, 3, 3] - Ks: Camera intrinsics. [..., C, 3, 3] - width: Image width. - height: Image height. - - Returns: - A tuple: - - - **Projected means**. [..., C, N, 2] - - **Projected covariances**. [..., C, N, 2, 2] - """ - warnings.warn( - "persp_proj is deprecated and will be removed in a future release. " - "Use proj with ortho=False instead.", - DeprecationWarning, - ) - return proj(means, covars, Ks, width, height, ortho=False) - - -def proj( - means: Tensor, # [..., C, N, 3] - covars: Tensor, # [..., C, N, 3, 3] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", -) -> Tuple[Tensor, Tensor]: - """Projection of Gaussians (perspective or orthographic). - - Args: - means: Gaussian means. [..., C, N, 3] - covars: Gaussian covariances. [..., C, N, 3, 3] - Ks: Camera intrinsics. [..., C, 3, 3] - width: Image width. - height: Image height. - - Returns: - A tuple: - - - **Projected means**. [..., C, N, 2] - - **Projected covariances**. [..., C, N, 2, 2] - """ - assert ( - camera_model != "ftheta" - ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - - batch_dims = means.shape[:-3] - C, N = means.shape[-3:-1] - assert means.shape == batch_dims + (C, N, 3), means.shape - assert covars.shape == batch_dims + (C, N, 3, 3), covars.shape - assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape - means = means.contiguous() - covars = covars.contiguous() - Ks = Ks.contiguous() - return _Proj.apply(means, covars, Ks, width, height, camera_model) - - -def fully_fused_projection( - means: Tensor, # [..., N, 3] - covars: Optional[Tensor], # [..., N, 6] or None - quats: Optional[Tensor], # [..., N, 4] or None - scales: Optional[Tensor], # [..., N, 3] or None - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - eps2d: float = 0.3, - near_plane: float = 0.01, - far_plane: float = 1e10, - radius_clip: float = 0.0, - packed: bool = False, - sparse_grad: bool = False, - calc_compensations: bool = False, - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - opacities: Optional[Tensor] = None, # [..., N] or None -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: - """Projects Gaussians to 2D. - - This function fuse the process of computing covariances - (:func:`quat_scale_to_covar_preci()`), transforming to camera space (:func:`world_to_cam()`), - and projection (:func:`proj()`). - - .. note:: - - During projection, we ignore the Gaussians that are outside of the camera frustum. - So not all the elements in the output tensors are valid. The output `radii` could serve as - 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 - 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 sycl kernel. Either `covars` or - {`quats`, `scales`} should be provided. - - Args: - means: Gaussian means. [..., N, 3] - covars: Gaussian covariances (flattened upper triangle). [..., N, 6] Optional. - quats: Quaternions (No need to be normalized). [..., N, 4] Optional. - scales: Scales. [..., N, 3] Optional. - viewmats: World-to-camera matrices. [..., C, 4, 4] - Ks: Camera intrinsics. [..., C, 3, 3] - width: Image width. - height: Image height. - eps2d: A epsilon added to the 2D covariance for numerical stability. Default: 0.3. - near_plane: Near plane distance. Default: 0.01. - far_plane: Far plane distance. Default: 1e10. - radius_clip: Gaussians with projected radii smaller than this value will be ignored. Default: 0.0. - packed: If True, the output tensors will be packed into a flattened tensor. Default: False. - sparse_grad: This is only effective when `packed` is True. If True, during backward the gradients - of {`means`, `covars`, `quats`, `scales`} will be a sparse Tensor in COO layout. Default: False. - calc_compensations: If True, a view-dependent opacity compensation factor will be computed, which - is useful for anti-aliasing. Default: False. - opacities: Gaussian opacities in range [0, 1]. If provided, will use it to compute a tighter bounds. - [..., N] or None. Default: None. - - Returns: - A tuple: - - If `packed` is True: - - - **batch_ids**. The batch indices of the projected Gaussians. Int32 tensor of shape [nnz]. - - **camera_ids**. The camera indices of the projected Gaussians. Int32 tensor of shape [nnz]. - - **gaussian_ids**. The column indices of the projected Gaussians. Int32 tensor of shape [nnz]. - - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [nnz, 2]. - - **means**. Projected Gaussian means in 2D. [nnz, 2] - - **depths**. The z-depth of the projected Gaussians. [nnz] - - **conics**. Inverse of the projected covariances. Return the flattend upper triangle with [nnz, 3] - - **compensations**. The view-dependent opacity compensation factor. [nnz] - - If `packed` is False: - - - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [..., C, N, 2]. - - **means**. Projected Gaussian means in 2D. [..., C, N, 2] - - **depths**. The z-depth of the projected Gaussians. [..., C, N] - - **conics**. Inverse of the projected covariances. Return the flattend upper triangle with [..., C, N, 3] - - **compensations**. The view-dependent opacity compensation factor. [..., C, N] - """ - batch_dims = means.shape[:-2] - N = means.shape[-2] - C = viewmats.shape[-3] - assert means.shape == batch_dims + (N, 3), means.shape - assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape - assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape - means = means.contiguous() - if covars is not None: - assert covars.shape == batch_dims + (N, 6), covars.shape - covars = covars.contiguous() - else: - assert quats is not None, "covars or quats is required" - assert scales is not None, "covars or scales is required" - assert quats.shape == batch_dims + (N, 4), quats.shape - assert scales.shape == batch_dims + (N, 3), scales.shape - quats = quats.contiguous() - scales = scales.contiguous() - if sparse_grad: - assert packed, "sparse_grad is only supported when packed is True" - assert batch_dims == (), "sparse_grad does not support batch dimensions" - if opacities is not None: - assert opacities.shape == batch_dims + (N,), opacities.shape - opacities = opacities.contiguous() - - assert ( - camera_model != "ftheta" - ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - - viewmats = viewmats.contiguous() - Ks = Ks.contiguous() - if packed: - return _FullyFusedProjectionPacked.apply( - means, - covars, - quats, - scales, - viewmats, - Ks, - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - sparse_grad, - calc_compensations, - camera_model, - opacities, - ) - else: - return _FullyFusedProjection.apply( - means, - covars, - quats, - scales, - viewmats, - Ks, - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - calc_compensations, - camera_model, - opacities, - ) - - -@torch.no_grad() -def isect_tiles( - means2d: Tensor, # [..., N, 2] or [nnz, 2] - radii: Tensor, # [..., N, 2] or [nnz, 2] - depths: Tensor, # [..., N] or [nnz] - tile_size: int, - tile_width: int, - tile_height: int, - sort: bool = True, - segmented: bool = False, - packed: bool = False, - n_images: Optional[int] = None, - image_ids: Optional[Tensor] = None, - gaussian_ids: Optional[Tensor] = None, -) -> Tuple[Tensor, Tensor, Tensor]: - """Maps projected Gaussians to intersecting tiles. - - Args: - means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. - radii: Maximum radii of the projected Gaussians. [..., N, 2] if packed is False, [nnz, 2] if packed is True. - depths: Z-depth of the projected Gaussians. [..., N] if packed is False, [nnz] if packed is True. - tile_size: Tile size. - tile_width: Tile width. - tile_height: Tile height. - sort: If True, the returned intersections will be sorted by the intersection ids. Default: True. - segmented: If True, segmented radix sort will be used to sort the intersections. Default: False. - packed: If True, the input tensors are packed. Default: False. - n_images: Number of images. Required if packed is True. - image_ids: The image indices of the projected Gaussians. Required if packed is True. - gaussian_ids: The column indices of the projected Gaussians. Required if packed is True. - - Returns: - A tuple: - - - **Tiles per Gaussian**. The number of tiles intersected by each Gaussian. - Int32 [..., N] if packed is False, Int32 [nnz] if packed is True. - - **Intersection ids**. Each id is an 64-bit integer with the following - information: image_id (Xc bits) | tile_id (Xt bits) | depth (32 bits). - Xc and Xt are the maximum number of bits required to represent the image and - tile ids, respectively. Int64 [n_isects] - - **Flatten ids**. The global flatten indices in [I * N] or [nnz] (packed). [n_isects] - """ - if packed: - nnz = means2d.size(0) - assert means2d.shape == (nnz, 2), means2d.shape - assert radii.shape == (nnz, 2), radii.shape - assert depths.shape == (nnz,), depths.shape - assert image_ids is not None, "image_ids is required if packed is True" - assert gaussian_ids is not None, "gaussian_ids is required if packed is True" - assert n_images is not None, "n_images is required if packed is True" - image_ids = image_ids.contiguous() - gaussian_ids = gaussian_ids.contiguous() - I = n_images - - else: - image_dims = means2d.shape[:-2] - I = math.prod(image_dims) - N = means2d.shape[-2] - assert means2d.shape == image_dims + (N, 2), means2d.shape - 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_sycl_func("intersect_tile")( - means2d.contiguous(), - radii.contiguous(), - depths.contiguous(), - image_ids, - gaussian_ids, - I, - tile_size, - tile_width, - tile_height, - sort, - segmented, - ) - return tiles_per_gauss, isect_ids, flatten_ids - - -@torch.no_grad() -def isect_offset_encode( - isect_ids: Tensor, - n_images: int, - tile_width: int, - tile_height: int, -) -> Tensor: - """Encodes intersection ids to offsets. - - Args: - isect_ids: Intersection ids. [n_isects] - n_images: Number of images. - tile_width: Tile width. - tile_height: Tile height. - - Returns: - Offsets. [I, tile_height, tile_width] - """ - return _make_lazy_sycl_func("intersect_offset")( - isect_ids.contiguous(), n_images, tile_width, tile_height - ) - - -def rasterize_to_pixels( - means2d: Tensor, # [..., N, 2] or [nnz, 2] - conics: Tensor, # [..., N, 3] or [nnz, 3] - colors: Tensor, # [..., N, channels] or [nnz, channels] - opacities: Tensor, # [..., N] or [nnz] - image_width: int, - image_height: int, - tile_size: int, - isect_offsets: Tensor, # [..., tile_height, tile_width] - flatten_ids: Tensor, # [n_isects] - backgrounds: Optional[Tensor] = None, # [..., channels] - masks: Optional[Tensor] = None, # [..., tile_height, tile_width] - packed: bool = False, - absgrad: bool = False, -) -> Tuple[Tensor, Tensor]: - """Rasterizes Gaussians to pixels. - - Args: - means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. - conics: Inverse of the projected covariances with only upper triangle values. [..., N, 3] if packed is False, [nnz, 3] if packed is True. - colors: Gaussian colors or ND features. [..., N, channels] if packed is False, [nnz, channels] if packed is True. - opacities: Gaussian opacities that support per-view values. [..., N] if packed is False, [nnz] if packed is True. - image_width: Image width. - image_height: Image height. - tile_size: Tile size. - isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] - flatten_ids: The global flatten indices in [I * N] or [nnz] from `isect_tiles()`. [n_isects] - backgrounds: Background colors. [..., channels]. Default: None. - masks: Optional tile mask to skip rendering GS to masked tiles. [..., tile_height, tile_width]. Default: None. - packed: If True, the input tensors are expected to be packed with shape [nnz, ...]. Default: False. - absgrad: If True, the backward pass will compute a `.absgrad` attribute for `means2d`. Default: False. - - Returns: - A tuple: - - - **Rendered colors**. [..., image_height, image_width, channels] - - **Rendered alphas**. [..., image_height, image_width, 1] - """ - - image_dims = means2d.shape[:-2] - channels = colors.shape[-1] - device = means2d.device - if packed: - nnz = means2d.size(0) - assert means2d.shape == (nnz, 2), means2d.shape - assert conics.shape == (nnz, 3), conics.shape - assert colors.shape[0] == nnz, colors.shape - assert opacities.shape == (nnz,), opacities.shape - else: - N = means2d.size(-2) - assert means2d.shape == image_dims + (N, 2), means2d.shape - assert conics.shape == image_dims + (N, 3), conics.shape - assert colors.shape == image_dims + (N, channels), colors.shape - assert opacities.shape == image_dims + (N,), opacities.shape - if backgrounds is not None: - assert backgrounds.shape == image_dims + (channels,), backgrounds.shape - backgrounds = backgrounds.contiguous() - if masks is not None: - assert masks.shape == isect_offsets.shape, masks.shape - masks = masks.contiguous() - - # Pad the channels to the nearest supported number if necessary - if channels > 513 or channels == 0: - # TODO: maybe worth to support zero channels? - raise ValueError(f"Unsupported number of color channels: {channels}") - if channels not in ( - 1, - 2, - 3, - 4, - 5, - 8, - 9, - 16, - 17, - 32, - 33, - 64, - 65, - 128, - 129, - 256, - 257, - 512, - 513, - ): - padded_channels = (1 << (channels - 1).bit_length()) - channels - colors = torch.cat( - [ - colors, - torch.zeros(*colors.shape[:-1], padded_channels, device=device), - ], - dim=-1, - ) - if backgrounds is not None: - backgrounds = torch.cat( - [ - backgrounds, - torch.zeros( - *backgrounds.shape[:-1], padded_channels, device=device - ), - ], - dim=-1, - ) - else: - padded_channels = 0 - - tile_height, tile_width = isect_offsets.shape[-2:] - assert ( - tile_height * tile_size >= image_height - ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" - assert ( - tile_width * tile_size >= image_width - ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - - render_colors, render_alphas = _RasterizeToPixels.apply( - means2d.contiguous(), - conics.contiguous(), - colors.contiguous(), - opacities.contiguous(), - backgrounds, - masks, - image_width, - image_height, - tile_size, - isect_offsets.contiguous(), - flatten_ids.contiguous(), - absgrad, - ) - - if padded_channels > 0: - render_colors = render_colors[..., :-padded_channels] - return render_colors, render_alphas - - -def rasterize_to_pixels_eval3d( - means: Tensor, # [..., N, 3] - quats: Tensor, # [..., N, 4] - scales: Tensor, # [..., N, 3] - colors: Tensor, # [..., C, N, channels] or [nnz, channels] - opacities: Tensor, # [..., C, N] or [nnz] - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - image_width: int, - image_height: int, - tile_size: int, - isect_offsets: Tensor, # [..., C, tile_height, tile_width] - flatten_ids: Tensor, # [n_isects] - backgrounds: Optional[Tensor] = None, # [..., C, channels] - masks: Optional[Tensor] = None, # [..., C, tile_height, tile_width] - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - ut_params: UnscentedTransformParameters = UnscentedTransformParameters(), - # distortion - 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, - # rolling shutter - rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, - viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] -) -> Tuple[Tensor, Tensor]: - """Rasterizes Gaussians to pixels. - - Similar to `rasterize_to_pixels()`, but compute the Gaussian responses in the - 3D world space instead of the 2D image space. Supports rolling shutter and - camera distortion. - - Returns: - A tuple: - - - **Rendered colors**. [..., C, image_height, image_width, channels] - - **Rendered alphas**. [..., C, image_height, image_width, 1] - """ - batch_dims = means.shape[:-2] - num_batch_dims = len(batch_dims) - N = means.size(-2) - C = viewmats.size(-3) - channels = colors.shape[-1] - device = means.device - - assert means.shape == batch_dims + (N, 3), means.shape - assert quats.shape == batch_dims + (N, 4), quats.shape - assert scales.shape == batch_dims + (N, 3), scales.shape - assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape - assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape - - assert colors.ndim in (num_batch_dims + 2, num_batch_dims + 3), colors.shape - if colors.ndim == num_batch_dims + 2: - raise NotImplementedError("packed mode is not supported yet") - assert ( - colors.shape[:-2] == batch_dims and colors.shape[-1] == channels - ), colors.shape - else: - assert colors.shape == batch_dims + (C, N, channels), colors.shape - assert opacities.shape == colors.shape[:-1], opacities.shape - - if backgrounds is not None: - assert backgrounds.shape == batch_dims + (C, channels), backgrounds.shape - backgrounds = backgrounds.contiguous() - - if masks is not None: - assert masks.shape == isect_offsets.shape, masks.shape - masks = masks.contiguous() - - if radial_coeffs is not None: - assert radial_coeffs.shape[:-1] == batch_dims + (C,) and radial_coeffs.shape[ - -1 - ] in (6, 4), radial_coeffs.shape - radial_coeffs = radial_coeffs.contiguous() - - if tangential_coeffs is not None: - assert tangential_coeffs.shape == batch_dims + (C, 2), tangential_coeffs.shape - tangential_coeffs = tangential_coeffs.contiguous() - - if thin_prism_coeffs is not None: - assert thin_prism_coeffs.shape == batch_dims + (C, 4), thin_prism_coeffs.shape - thin_prism_coeffs = thin_prism_coeffs.contiguous() - - if viewmats_rs is not None: - assert viewmats_rs.shape == batch_dims + (C, 4, 4), viewmats_rs.shape - viewmats_rs = viewmats_rs.contiguous() - - # Pad the channels to the nearest supported number if necessary - channels = colors.shape[-1] - if channels > 513 or channels == 0: - # TODO: maybe worth to support zero channels? - raise ValueError(f"Unsupported number of color channels: {channels}") - if channels not in ( - 1, - 2, - 3, - 4, - 5, - 8, - 9, - 16, - 17, - 32, - 33, - 64, - 65, - 128, - 129, - 256, - 257, - 512, - 513, - ): - padded_channels = (1 << (channels - 1).bit_length()) - channels - colors = torch.cat( - [ - colors, - torch.zeros(*colors.shape[:-1], padded_channels, device=device), - ], - dim=-1, - ) - if backgrounds is not None: - backgrounds = torch.cat( - [ - backgrounds, - torch.zeros( - *backgrounds.shape[:-1], padded_channels, device=device - ), - ], - dim=-1, - ) - else: - padded_channels = 0 - - tile_height, tile_width = isect_offsets.shape[-2:] - assert ( - tile_height * tile_size >= image_height - ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" - assert ( - tile_width * tile_size >= image_width - ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - - render_colors, render_alphas = _RasterizeToPixelsEval3D.apply( - means.contiguous(), - quats.contiguous(), - scales.contiguous(), - colors.contiguous(), - opacities.contiguous(), - backgrounds.contiguous() if backgrounds is not None else None, - masks.contiguous() if masks is not None else None, - viewmats.contiguous(), - Ks.contiguous(), - image_width, - image_height, - tile_size, - isect_offsets.contiguous(), - flatten_ids.contiguous(), - camera_model, - ut_params, - # distortion - 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, - # rolling shutter - rolling_shutter, - viewmats_rs.contiguous() if viewmats_rs is not None else None, - ) - - if padded_channels > 0: - render_colors = render_colors[..., :-padded_channels] - return render_colors, render_alphas - - -@torch.no_grad() -def rasterize_to_indices_in_range( - range_start: int, - range_end: int, - transmittances: Tensor, # [..., image_height, image_width] - means2d: Tensor, # [..., N, 2] - conics: Tensor, # [..., N, 3] - opacities: Tensor, # [..., N] - image_width: int, - image_height: int, - tile_size: int, - isect_offsets: Tensor, # [..., tile_height, tile_width] - flatten_ids: Tensor, # [n_isects] -) -> Tuple[Tensor, Tensor, Tensor]: - """Rasterizes a batch of Gaussians to images but only returns the indices. - - .. note:: - - This function supports iterative rasterization, in which each call of this function - will rasterize a batch of Gaussians from near to far, defined by `[range_start, range_end)`. - If a one-step full rasterization is desired, set `range_start` to 0 and `range_end` to a really - large number, e.g, 1e10. - - Args: - range_start: The start batch of Gaussians to be rasterized (inclusive). - range_end: The end batch of Gaussians to be rasterized (exclusive). - transmittances: Currently transmittances. [..., image_height, image_width] - means2d: Projected Gaussian means. [..., N, 2] - conics: Inverse of the projected covariances with only upper triangle values. [..., N, 3] - opacities: Gaussian opacities that support per-view values. [..., N] - image_width: Image width. - image_height: Image height. - tile_size: Tile size. - isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] - flatten_ids: The global flatten indices in [I * N] from `isect_tiles()`. [n_isects] - - Returns: - A tuple: - - - **Gaussian ids**. Gaussian ids for the pixel intersection. A flattened list of shape [M]. - - **Pixel ids**. pixel indices (row-major). A flattened list of shape [M]. - - **Image ids**. image indices. A flattened list of shape [M]. - """ - - image_dims = means2d.shape[:-2] - tile_height, tile_width = isect_offsets.shape[-2:] - N = means2d.shape[-2] - assert transmittances.shape == image_dims + ( - image_height, - image_width, - ), transmittances.shape - assert means2d.shape == image_dims + (N, 2), means2d.shape - assert conics.shape == image_dims + (N, 3), conics.shape - assert opacities.shape == image_dims + (N,), opacities.shape - assert isect_offsets.shape == image_dims + ( - tile_height, - tile_width, - ), isect_offsets.shape - assert ( - tile_height * tile_size >= image_height - ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" - assert ( - tile_width * tile_size >= image_width - ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - - out_gauss_ids, out_indices = _make_lazy_sycl_func("rasterize_to_indices_3dgs")( - range_start, - range_end, - transmittances.contiguous(), - means2d.contiguous(), - conics.contiguous(), - opacities.contiguous(), - image_width, - image_height, - tile_size, - isect_offsets.contiguous(), - flatten_ids.contiguous(), - ) - out_pixel_ids = out_indices % (image_width * image_height) - out_image_ids = out_indices // (image_width * image_height) - return out_gauss_ids, out_pixel_ids, out_image_ids - - -class _QuatScaleToCovarPreci(torch.autograd.Function): - """Converts quaternions and scales to covariance and precision matrices.""" - - @staticmethod - def forward( - ctx, - quats: Tensor, # [..., 4], - scales: Tensor, # [..., 3], - compute_covar: bool = True, - compute_preci: bool = True, - triu: bool = False, - ) -> Tuple[Tensor, Tensor]: - covars, precis = _make_lazy_sycl_func("quat_scale_to_covar_preci_fwd")( - quats, scales, compute_covar, compute_preci, triu - ) - ctx.save_for_backward(quats, scales) - ctx.compute_covar = compute_covar - ctx.compute_preci = compute_preci - ctx.triu = triu - return covars, precis - - @staticmethod - def backward(ctx, v_covars: Tensor, v_precis: Tensor): - quats, scales = ctx.saved_tensors - compute_covar = ctx.compute_covar - compute_preci = ctx.compute_preci - triu = ctx.triu - if compute_covar and v_covars.is_sparse: - 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_sycl_func("quat_scale_to_covar_preci_bwd")( - quats, - scales, - triu, - v_covars.contiguous() if compute_covar else None, - v_precis.contiguous() if compute_preci else None, - ) - return v_quats, v_scales, None, None, None - - -class _Proj(torch.autograd.Function): - """Perspective fully_fused_projection on Gaussians.""" - - @staticmethod - def forward( - ctx, - means: Tensor, # [..., C, N, 3] - covars: Tensor, # [..., C, N, 3, 3] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - ) -> Tuple[Tensor, Tensor]: - assert ( - camera_model != "ftheta" - ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - - camera_model_type = _make_lazy_sycl_obj( - f"CameraModelType.{camera_model.upper()}" - ) - - means2d, covars2d = _make_lazy_sycl_func("projection_ewa_simple_fwd")( - means, - covars, - Ks, - width, - height, - camera_model_type, - ) - ctx.save_for_backward(means, covars, Ks) - ctx.width = width - ctx.height = height - ctx.camera_model_type = camera_model_type - return means2d, covars2d - - @staticmethod - def backward(ctx, v_means2d: Tensor, v_covars2d: Tensor): - means, covars, Ks = ctx.saved_tensors - width = ctx.width - height = ctx.height - camera_model_type = ctx.camera_model_type - v_means, v_covars = _make_lazy_sycl_func("projection_ewa_simple_bwd")( - means, - covars, - Ks, - width, - height, - camera_model_type, - v_means2d.contiguous(), - v_covars2d.contiguous(), - ) - return v_means, v_covars, None, None, None, None - - -class _FullyFusedProjection(torch.autograd.Function): - """Projects Gaussians to 2D.""" - - @staticmethod - def forward( - ctx, - means: Tensor, # [..., N, 3] - covars: Tensor, # [..., N, 6] or None - quats: Tensor, # [..., N, 4] or None - scales: Tensor, # [..., N, 3] or None - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - eps2d: float, - near_plane: float, - far_plane: float, - radius_clip: float, - calc_compensations: bool, - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - opacities: Optional[Tensor] = None, # [..., N] or None - ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: - assert ( - camera_model != "ftheta" - ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - - camera_model_type = _make_lazy_sycl_obj( - f"CameraModelType.{camera_model.upper()}" - ) - - # "covars" and {"quats", "scales"} are mutually exclusive - radii, means2d, depths, conics, compensations = _make_lazy_sycl_func( - "projection_ewa_3dgs_fused_fwd" - )( - means, - covars, - quats, - scales, - opacities, - viewmats, - Ks, - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - calc_compensations, - camera_model_type, - ) - if not calc_compensations: - compensations = None - ctx.save_for_backward( - means, covars, quats, scales, viewmats, Ks, radii, conics, compensations - ) - ctx.width = width - ctx.height = height - ctx.eps2d = eps2d - ctx.camera_model_type = camera_model_type - - return radii, means2d, depths, conics, compensations - - @staticmethod - def backward(ctx, v_radii, v_means2d, v_depths, v_conics, v_compensations): - ( - means, - covars, - quats, - scales, - viewmats, - Ks, - radii, - conics, - compensations, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - eps2d = ctx.eps2d - 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_sycl_func( - "projection_ewa_3dgs_fused_bwd" - )( - means, - covars, - quats, - scales, - viewmats, - Ks, - width, - height, - eps2d, - camera_model_type, - radii, - conics, - compensations, - v_means2d.contiguous(), - v_depths.contiguous(), - v_conics.contiguous(), - v_compensations, - ctx.needs_input_grad[4], # viewmats_requires_grad - ) - if not ctx.needs_input_grad[0]: - v_means = None - if not ctx.needs_input_grad[1]: - v_covars = None - if not ctx.needs_input_grad[2]: - v_quats = None - if not ctx.needs_input_grad[3]: - v_scales = None - if not ctx.needs_input_grad[4]: - v_viewmats = None - return ( - v_means, - v_covars, - v_quats, - v_scales, - v_viewmats, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def fully_fused_projection_with_ut( - means: Tensor, # [..., N, 3] - quats: Tensor, # [..., N, 4] - scales: Tensor, # [..., N, 3] - opacities: Optional[Tensor], # [..., N] - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - eps2d: float = 0.3, - near_plane: float = 0.01, - far_plane: float = 1e10, - radius_clip: float = 0.0, - calc_compensations: bool = False, - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - ut_params: UnscentedTransformParameters = UnscentedTransformParameters(), - # distortion - 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, - # rolling shutter - rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, - viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] -) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: - """Projects Gaussians to 2D using Unscented Transform (UT). - - similar to `fully_fused_projection()`, but supports camera distortion and - rolling shutter. - - .. warning:: - This function is not differentiable to any input. - """ - batch_dims = means.shape[:-2] - N = means.shape[-2] - C = viewmats.shape[-3] - assert means.shape == batch_dims + (N, 3), means.shape - assert quats.shape == batch_dims + (N, 4), quats.shape - assert scales.shape == batch_dims + (N, 3), scales.shape - if opacities is not None: - assert opacities.shape == batch_dims + (N,), opacities.shape - assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape - assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape - if radial_coeffs is not None: - assert radial_coeffs.shape[:-1] == batch_dims + (C,) and radial_coeffs.shape[ - -1 - ] in [6, 4], radial_coeffs.shape - if tangential_coeffs is not None: - assert tangential_coeffs.shape == batch_dims + (C, 2), tangential_coeffs.shape - if thin_prism_coeffs is not None: - assert thin_prism_coeffs.shape == batch_dims + (C, 4), thin_prism_coeffs.shape - if viewmats_rs is not None: - assert viewmats_rs.shape == batch_dims + (C, 4, 4), viewmats_rs.shape - - camera_model_type = _make_lazy_sycl_obj(f"CameraModelType.{camera_model.upper()}") - - radii, means2d, depths, conics, compensations = _make_lazy_sycl_func( - "projection_ut_3dgs_fused" - )( - means.contiguous(), - quats.contiguous(), - scales.contiguous(), - opacities.contiguous() if opacities is not None else None, - viewmats.contiguous(), - viewmats_rs.contiguous() if viewmats_rs is not None else None, - Ks.contiguous(), - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - calc_compensations, - camera_model_type, - ut_params.to_cpp(), - rolling_shutter.to_cpp(), - 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() - ), - ) - if not calc_compensations: - compensations = None - return radii, means2d, depths, conics, compensations - - -class _RasterizeToPixels(torch.autograd.Function): - """Rasterize gaussians""" - - @staticmethod - def forward( - ctx, - means2d: Tensor, # [..., N, 2] or [nnz, 2] - conics: Tensor, # [..., N, 3] or [nnz, 3] - colors: Tensor, # [..., N, channels] or [nnz, channels] - opacities: Tensor, # [..., N] or [nnz] - backgrounds: Tensor, # [..., channels], Optional - masks: Tensor, # [..., tile_height, tile_width], Optional - width: int, - height: int, - tile_size: int, - isect_offsets: Tensor, # [..., tile_height, tile_width] - flatten_ids: Tensor, # [n_isects] - absgrad: bool, - ) -> Tuple[Tensor, Tensor]: - render_colors, render_alphas, last_ids = _make_lazy_sycl_func( - "rasterize_to_pixels_3dgs_fwd" - )( - means2d, - conics, - colors, - opacities, - backgrounds, - masks, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - ) - - ctx.save_for_backward( - means2d, - conics, - colors, - opacities, - backgrounds, - masks, - isect_offsets, - flatten_ids, - render_alphas, - last_ids, - ) - ctx.width = width - ctx.height = height - ctx.tile_size = tile_size - ctx.absgrad = absgrad - - # double to float - render_alphas = render_alphas.float() - return render_colors, render_alphas - - @staticmethod - def backward( - ctx, - v_render_colors: Tensor, # [..., H, W, 3] - v_render_alphas: Tensor, # [..., H, W, 1] - ): - ( - means2d, - conics, - colors, - opacities, - backgrounds, - masks, - isect_offsets, - flatten_ids, - render_alphas, - last_ids, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - tile_size = ctx.tile_size - absgrad = ctx.absgrad - - ( - v_means2d_abs, - v_means2d, - v_conics, - v_colors, - v_opacities, - ) = _make_lazy_sycl_func("rasterize_to_pixels_3dgs_bwd")( - means2d, - conics, - colors, - opacities, - backgrounds, - masks, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - render_alphas, - last_ids, - v_render_colors.contiguous(), - v_render_alphas.contiguous(), - absgrad, - ) - - if absgrad: - means2d.absgrad = v_means2d_abs - - if ctx.needs_input_grad[4]: - v_backgrounds = (v_render_colors * (1.0 - render_alphas).float()).sum( - dim=(-3, -2) - ) - else: - v_backgrounds = None - - return ( - v_means2d, - v_conics, - v_colors, - v_opacities, - v_backgrounds, - None, - None, - None, - None, - None, - None, - None, - ) - - -class _RasterizeToPixelsEval3D(torch.autograd.Function): - """Rasterize gaussians""" - - @staticmethod - def forward( - ctx, - means: Tensor, # [..., N, 3] - quats: Tensor, # [..., N, 4] - scales: Tensor, # [..., N, 3] - colors: Tensor, # [..., C, N, D] or [nnz, D] - opacities: Tensor, # [..., C, N] or [nnz] - backgrounds: Tensor, # [..., C, D], Optional - masks: Tensor, # [..., C, tile_height, tile_width], Optional - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - tile_size: int, - isect_offsets: Tensor, # [..., C, tile_height, tile_width] - flatten_ids: Tensor, # [..., n_isects] - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - ut_params: UnscentedTransformParameters = UnscentedTransformParameters(), - # distortion - 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, - # rolling shutter - rolling_shutter: RollingShutterType = RollingShutterType.GLOBAL, - viewmats_rs: Optional[Tensor] = None, # [..., C, 4, 4] - ) -> Tuple[Tensor, Tensor]: - ut_params = ut_params.to_cpp() - rs_type = rolling_shutter.to_cpp() - camera_model_type = _make_lazy_sycl_obj( - f"CameraModelType.{camera_model.upper()}" - ) - ftheta_coeffs = ( - ftheta_coeffs.to_cpp() - if ftheta_coeffs is not None - else FThetaCameraDistortionParameters.to_cpp_default() - ) - - render_colors, render_alphas, last_ids = _make_lazy_sycl_func( - "rasterize_to_pixels_from_world_3dgs_fwd" - )( - means, - quats, - scales, - colors, - opacities, - backgrounds, - masks, - width, - height, - tile_size, - viewmats, - viewmats_rs, - Ks, - camera_model_type, - ut_params, - rs_type, - radial_coeffs, - tangential_coeffs, - thin_prism_coeffs, - ftheta_coeffs, - isect_offsets, - flatten_ids, - ) - - ctx.save_for_backward( - means, - quats, - scales, - colors, - opacities, - backgrounds, - masks, - viewmats, - viewmats_rs, - Ks, - radial_coeffs, - tangential_coeffs, - thin_prism_coeffs, - isect_offsets, - flatten_ids, - render_alphas, - last_ids, - ) - ctx.width = width - ctx.height = height - ctx.ut_params = ut_params - ctx.rs_type = rs_type - ctx.camera_model_type = camera_model_type - ctx.tile_size = tile_size - ctx.ftheta_coeffs = ftheta_coeffs - - return render_colors, render_alphas - - @staticmethod - def backward( - ctx, - v_render_colors: Tensor, # [..., C, H, W, 3] - v_render_alphas: Tensor, # [..., C, H, W, 1] - ): - ( - means, - quats, - scales, - colors, - opacities, - backgrounds, - masks, - viewmats, - viewmats_rs, - Ks, - radial_coeffs, - tangential_coeffs, - thin_prism_coeffs, - isect_offsets, - flatten_ids, - render_alphas, - last_ids, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - ut_params = ctx.ut_params - rs_type = ctx.rs_type - camera_model_type = ctx.camera_model_type - tile_size = ctx.tile_size - ftheta_coeffs = ctx.ftheta_coeffs - - (v_means, v_quats, v_scales, v_colors, v_opacities,) = _make_lazy_sycl_func( - "rasterize_to_pixels_from_world_3dgs_bwd" - )( - means, - quats, - scales, - colors, - opacities, - backgrounds, - masks, - width, - height, - tile_size, - viewmats, - viewmats_rs, - Ks, - camera_model_type, - ut_params, - rs_type, - radial_coeffs, - tangential_coeffs, - thin_prism_coeffs, - ftheta_coeffs, - isect_offsets, - flatten_ids, - render_alphas, - last_ids, - v_render_colors.contiguous(), - v_render_alphas.contiguous(), - ) - - if ctx.needs_input_grad[5]: # backgrounds - v_backgrounds = (v_render_colors * (1.0 - render_alphas).float()).sum( - dim=(-3, -2) - ) - else: - v_backgrounds = None - - if ctx.needs_input_grad[7]: # viewmats - raise NotImplementedError - - return ( - v_means, - v_quats, - v_scales, - v_colors, - v_opacities, - v_backgrounds, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class _FullyFusedProjectionPacked(torch.autograd.Function): - """Projects Gaussians to 2D. Return packed tensors.""" - - @staticmethod - def forward( - ctx, - means: Tensor, # [..., N, 3] - covars: Tensor, # [..., N, 6] or None - quats: Tensor, # [..., N, 4] or None - scales: Tensor, # [..., N, 3] or None - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - eps2d: float, - near_plane: float, - far_plane: float, - radius_clip: float, - sparse_grad: bool, - calc_compensations: bool, - camera_model: Literal["pinhole", "ortho", "fisheye", "ftheta"] = "pinhole", - opacities: Optional[Tensor] = None, # [..., N] or None - ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - assert ( - camera_model != "ftheta" - ), "ftheta camera is only supported via UT, please set with_ut=True in the rasterization()" - - camera_model_type = _make_lazy_sycl_obj( - f"CameraModelType.{camera_model.upper()}" - ) - - ( - batch_ids, - camera_ids, - gaussian_ids, - radii, - means2d, - depths, - conics, - indptr, - compensations, - ) = _make_lazy_sycl_func("projection_ewa_3dgs_packed_fwd")( - means, - covars, # optional - quats, # optional - scales, # optional - opacities, # optional - viewmats, - Ks, - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - calc_compensations, - camera_model_type, - ) - if not calc_compensations: - compensations = None - ctx.save_for_backward( - batch_ids, - camera_ids, - gaussian_ids, - means, - covars, - quats, - scales, - viewmats, - Ks, - conics, - compensations, - ) - ctx.width = width - ctx.height = height - ctx.eps2d = eps2d - ctx.sparse_grad = sparse_grad - ctx.camera_model_type = camera_model_type - - return ( - batch_ids, - camera_ids, - gaussian_ids, - radii, - means2d, - depths, - conics, - compensations, - ) - - @staticmethod - def backward( - ctx, - v_batch_ids, - v_camera_ids, - v_gaussian_ids, - v_radii, - v_means2d, - v_depths, - v_conics, - v_compensations, - ): - ( - batch_ids, - camera_ids, - gaussian_ids, - means, - covars, - quats, - scales, - viewmats, - Ks, - conics, - compensations, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - eps2d = ctx.eps2d - sparse_grad = ctx.sparse_grad - 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_sycl_func( - "projection_ewa_3dgs_packed_bwd" - )( - means, - covars, - quats, - scales, - viewmats, - Ks, - width, - height, - eps2d, - camera_model_type, - batch_ids, - camera_ids, - gaussian_ids, - conics, - compensations, - v_means2d.contiguous(), - v_depths.contiguous(), - v_conics.contiguous(), - v_compensations, - ctx.needs_input_grad[4], # viewmats_requires_grad - sparse_grad, - ) - - if sparse_grad: - batch_dims = means.shape[:-2] - B = math.prod(batch_dims) - N = means.shape[-2] - if not ctx.needs_input_grad[0]: - v_means = None - else: - if sparse_grad: - # TODO: gaussian_ids is duplicated so not ideal. - # An idea is to directly set the attribute (e.g., .sparse_grad) of - # the tensor but this requires the tensor to be leaf node only. And - # a customized optimizer would be needed in this case. - v_means = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_means, # [nnz, 3] - size=means.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[1]: - v_covars = None - else: - if sparse_grad: - v_covars = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_covars, # [nnz, 6] - size=covars.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[2]: - v_quats = None - else: - if sparse_grad: - v_quats = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_quats, # [nnz, 4] - size=quats.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[3]: - v_scales = None - else: - if sparse_grad: - v_scales = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_scales, # [nnz, 3] - size=scales.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[4]: - v_viewmats = None - - return ( - v_means, - v_covars, - v_quats, - v_scales, - v_viewmats, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class _SphericalHarmonics(torch.autograd.Function): - """Spherical Harmonics""" - - @staticmethod - def forward( - ctx, sh_degree: int, dirs: Tensor, coeffs: Tensor, masks: Tensor - ) -> Tensor: - colors = _make_lazy_sycl_func("spherical_harmonics_fwd")( - sh_degree, dirs, coeffs, masks - ) - ctx.save_for_backward(dirs, coeffs, masks) - ctx.sh_degree = sh_degree - ctx.num_bases = coeffs.shape[-2] - return colors - - @staticmethod - def backward(ctx, v_colors: Tensor): - dirs, coeffs, masks = ctx.saved_tensors - sh_degree = ctx.sh_degree - num_bases = ctx.num_bases - compute_v_dirs = ctx.needs_input_grad[1] - v_coeffs, v_dirs = _make_lazy_sycl_func("spherical_harmonics_bwd")( - num_bases, - sh_degree, - dirs, - coeffs, - masks, - v_colors.contiguous(), - compute_v_dirs, - ) - if not compute_v_dirs: - v_dirs = None - return None, v_dirs, v_coeffs, None - - -###### 2DGS ###### -def fully_fused_projection_2dgs( - means: Tensor, # [..., N, 3] - quats: Tensor, # [..., N, 4] - scales: Tensor, # [..., N, 3] - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - eps2d: float = 0.3, - near_plane: float = 0.01, - far_plane: float = 1e10, - radius_clip: float = 0.0, - packed: bool = False, - sparse_grad: bool = False, -) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - """Prepare Gaussians for rasterization - - This function prepares ray-splat intersection matrices, computes - per splat bounding box and 2D means in image space. - - Args: - means: Gaussian means. [..., N, 3] - quats: Quaternions (No need to be normalized). [..., N, 4]. - scales: Scales. [..., N, 3]. - viewmats: World-to-camera matrices. [..., C, 4, 4] - Ks: Camera intrinsics. [..., C, 3, 3] - width: Image width. - height: Image height. - near_plane: Near plane distance. Default: 0.01. - far_plane: Far plane distance. Default: 200. - radius_clip: Gaussians with projected radii smaller than this value will be ignored. Default: 0.0. - packed: If True, the output tensors will be packed into a flattened tensor. Default: False. - sparse_grad (Experimental): This is only effective when `packed` is True. If True, during backward the gradients - of {`means`, `covars`, `quats`, `scales`} will be a sparse Tensor in COO layout. Default: False. - - Returns: - A tuple: - - If `packed` is True: - - - **batch_ids**. The batch indices of the projected Gaussians. Int32 tensor of shape [nnz]. - - **camera_ids**. The camera indices of the projected Gaussians. Int32 tensor of shape [nnz]. - - **gaussian_ids**. The column indices of the projected Gaussians. Int32 tensor of shape [nnz]. - - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [nnz, 2]. - - **means**. Projected Gaussian means in 2D. [nnz, 2] - - **depths**. The z-depth of the projected Gaussians. [nnz] - - **ray_transforms**. transformation matrices that transforms xy-planes in pixel spaces into splat coordinates (WH)^T in equation (9) in paper [nnz, 3, 3] - - **normals**. The normals in camera spaces. [nnz, 3] - - If `packed` is False: - - - **radii**. The maximum radius of the projected Gaussians in pixel unit. Int32 tensor of shape [..., C, N, 2]. - - **means**. Projected Gaussian means in 2D. [..., C, N, 2] - - **depths**. The z-depth of the projected Gaussians. [..., C, N] - - **ray_transforms**. transformation matrices that transforms xy-planes in pixel spaces into splat coordinates [..., C, N, 3, 3] - - **normals**. The normals in camera spaces. [..., C, N, 3] - - """ - batch_dims = means.shape[:-2] - N = means.shape[-2] - C = viewmats.shape[-3] - assert means.shape == batch_dims + (N, 3), means.shape - assert viewmats.shape == batch_dims + (C, 4, 4), viewmats.shape - assert Ks.shape == batch_dims + (C, 3, 3), Ks.shape - means = means.contiguous() - assert quats is not None, "quats is required" - assert scales is not None, "scales is required" - assert quats.shape == batch_dims + (N, 4), quats.shape - assert scales.shape == batch_dims + (N, 3), scales.shape - quats = quats.contiguous() - scales = scales.contiguous() - if sparse_grad: - assert packed, "sparse_grad is only supported when packed is True" - - viewmats = viewmats.contiguous() - Ks = Ks.contiguous() - if packed: - return _FullyFusedProjectionPacked2DGS.apply( - means, - quats, - scales, - viewmats, - Ks, - width, - height, - near_plane, - far_plane, - radius_clip, - sparse_grad, - ) - else: - return _FullyFusedProjection2DGS.apply( - means, - quats, - scales, - viewmats, - Ks, - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - ) - - -class _FullyFusedProjection2DGS(torch.autograd.Function): - """Projects Gaussians to 2D.""" - - @staticmethod - def forward( - ctx, - means: Tensor, # [..., N, 3] - quats: Tensor, # [..., N, 4] - scales: Tensor, # [..., N, 3] - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - eps2d: float, - near_plane: float, - far_plane: float, - radius_clip: float, - ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - radii, means2d, depths, ray_transforms, normals = _make_lazy_sycl_func( - "projection_2dgs_fused_fwd" - )( - means, - quats, - scales, - viewmats, - Ks, - width, - height, - eps2d, - near_plane, - far_plane, - radius_clip, - ) - ctx.save_for_backward( - means, - quats, - scales, - viewmats, - Ks, - radii, - ray_transforms, - normals, - ) - ctx.width = width - ctx.height = height - ctx.eps2d = eps2d - - return radii, means2d, depths, ray_transforms, normals - - @staticmethod - def backward(ctx, v_radii, v_means2d, v_depths, v_ray_transforms, v_normals): - ( - means, - quats, - scales, - viewmats, - Ks, - radii, - ray_transforms, - normals, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - eps2d = ctx.eps2d - v_means, v_quats, v_scales, v_viewmats = _make_lazy_sycl_func( - "projection_2dgs_fused_bwd" - )( - means, - quats, - scales, - viewmats, - Ks, - width, - height, - radii, - ray_transforms, - v_means2d.contiguous(), - v_depths.contiguous(), - v_normals.contiguous(), - v_ray_transforms.contiguous(), - ctx.needs_input_grad[3], # viewmats_requires_grad - ) - if not ctx.needs_input_grad[0]: - v_means = None - if not ctx.needs_input_grad[1]: - v_quats = None - if not ctx.needs_input_grad[2]: - v_scales = None - if not ctx.needs_input_grad[3]: - v_viewmats = None - - return ( - v_means, - v_quats, - v_scales, - v_viewmats, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -class _FullyFusedProjectionPacked2DGS(torch.autograd.Function): - """Projects Gaussians to 2D. Return packed tensors.""" - - @staticmethod - def forward( - ctx, - means: Tensor, # [..., N, 3] - quats: Tensor, # [..., N, 4] - scales: Tensor, # [..., N, 3] - viewmats: Tensor, # [..., C, 4, 4] - Ks: Tensor, # [..., C, 3, 3] - width: int, - height: int, - near_plane: float, - far_plane: float, - radius_clip: float, - sparse_grad: bool, - ) -> Tuple[Tensor, Tensor, Tensor, Tensor]: - ( - indptr, - batch_ids, - camera_ids, - gaussian_ids, - radii, - means2d, - depths, - ray_transforms, - normals, - ) = _make_lazy_sycl_func("projection_2dgs_packed_fwd")( - means, - quats, - scales, - viewmats, - Ks, - width, - height, - near_plane, - far_plane, - radius_clip, - ) - ctx.save_for_backward( - batch_ids, - camera_ids, - gaussian_ids, - means, - quats, - scales, - viewmats, - Ks, - ray_transforms, - ) - ctx.width = width - ctx.height = height - ctx.sparse_grad = sparse_grad - - return ( - batch_ids, - camera_ids, - gaussian_ids, - radii, - means2d, - depths, - ray_transforms, - normals, - ) - - @staticmethod - def backward( - ctx, - v_batch_ids, - v_camera_ids, - v_gaussian_ids, - v_radii, - v_means2d, - v_depths, - v_ray_transforms, - v_normals, - ): - ( - batch_ids, - camera_ids, - gaussian_ids, - means, - quats, - scales, - viewmats, - Ks, - ray_transforms, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - sparse_grad = ctx.sparse_grad - - v_means, v_quats, v_scales, v_viewmats = _make_lazy_sycl_func( - "projection_2dgs_packed_bwd" - )( - means, - quats, - scales, - viewmats, - Ks, - width, - height, - batch_ids, - camera_ids, - gaussian_ids, - ray_transforms, - v_means2d.contiguous(), - v_depths.contiguous(), - v_ray_transforms.contiguous(), - v_normals.contiguous(), - ctx.needs_input_grad[3], # viewmats_requires_grad - sparse_grad, - ) - - if sparse_grad: - batch_dims = means.shape[:-2] - B = math.prod(batch_dims) - N = means.shape[-2] - - if not ctx.needs_input_grad[0]: - v_means = None - else: - if sparse_grad: - # TODO: gaussian_ids is duplicated so not ideal. - # An idea is to directly set the attribute (e.g., .sparse_grad) of - # the tensor but this requires the tensor to be leaf node only. And - # a customized optimizer would be needed in this case. - v_means = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_means, # [nnz, 3] - size=means.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[1]: - v_quats = None - else: - if sparse_grad: - v_quats = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_quats, # [nnz, 4] - size=quats.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[2]: - v_scales = None - else: - if sparse_grad: - v_scales = torch.sparse_coo_tensor( - indices=gaussian_ids[None], - values=v_scales, # [nnz, 3] - size=scales.shape, - is_coalesced=len(viewmats) == 1, - ) - if not ctx.needs_input_grad[3]: - v_viewmats = None - - return ( - v_means, - v_quats, - v_scales, - v_viewmats, - None, - None, - None, - None, - None, - None, - None, - None, - None, - ) - - -def rasterize_to_pixels_2dgs( - means2d: Tensor, # [..., N, 2] - ray_transforms: Tensor, # [..., N, 3, 3] - colors: Tensor, # [..., N, channels] - opacities: Tensor, # [..., N] - normals: Tensor, # [..., N, 3] - densify: Tensor, # [..., N, 2] - image_width: int, - image_height: int, - tile_size: int, - isect_offsets: Tensor, # [..., tile_height, tile_width] - flatten_ids: Tensor, # [n_isects] - backgrounds: Optional[Tensor] = None, # [..., channels] - masks: Optional[Tensor] = None, # [..., tile_height, tile_width] - packed: bool = False, - absgrad: bool = False, - distloss: bool = False, -) -> Tuple[Tensor, Tensor]: - """Rasterize Gaussians to pixels. - - Args: - means2d: Projected Gaussian means. [..., N, 2] if packed is False, [nnz, 2] if packed is True. - ray_transforms: transformation matrices that transforms xy-planes in pixel spaces into splat coordinates. [..., N, 3, 3] if packed is False, [nnz, channels] if packed is True. - colors: Gaussian colors or ND features. [..., N, channels] if packed is False, [nnz, channels] if packed is True. - opacities: Gaussian opacities that support per-view values. [..., N] if packed is False, [nnz] if packed is True. - normals: The normals in camera space. [..., N, 3] if packed is False, [nnz, 3] if packed is True. - densify: Dummy variable to keep track of gradient for densification. [..., N, 2] if packed, [nnz, 3] if packed is True. - tile_size: Tile size. - isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] - flatten_ids: The global flatten indices in [I * N] or [nnz] from `isect_tiles()`. [n_isects] - backgrounds: Background colors. [..., channels]. Default: None. - masks: Optional tile mask to skip rendering GS to masked tiles. [..., tile_height, tile_width]. Default: None. - packed: If True, the input tensors are expected to be packed with shape [nnz, ...]. Default: False. - absgrad: If True, the backward pass will compute a `.absgrad` attribute for `means2d`. Default: False. - - Returns: - A tuple: - - - **Rendered colors**. [..., image_height, image_width, channels] - - **Rendered alphas**. [..., image_height, image_width, 1] - - **Rendered normals**. [..., image_height, image_width, 3] - - **Rendered distortion**. [..., image_height, image_width, 1] - - **Rendered median depth**.[..., image_height, image_width, 1] - - - """ - image_dims = means2d.shape[:-2] - channels = colors.shape[-1] - device = means2d.device - if packed: - nnz = means2d.size(0) - assert means2d.shape == (nnz, 2), means2d.shape - assert ray_transforms.shape == (nnz, 3, 3), ray_transforms.shape - assert colors.shape[0] == nnz, colors.shape - assert opacities.shape == (nnz,), opacities.shape - else: - N = means2d.size(-2) - assert means2d.shape == image_dims + (N, 2), means2d.shape - assert ray_transforms.shape == image_dims + (N, 3, 3), ray_transforms.shape - assert colors.shape[:-2] == image_dims, colors.shape - assert opacities.shape == image_dims + (N,), opacities.shape - if backgrounds is not None: - assert backgrounds.shape == image_dims + (channels,), backgrounds.shape - backgrounds = backgrounds.contiguous() - - # Pad the channels to the nearest supported number if necessary - if channels > 512 or channels == 0: - # TODO: maybe worth to support zero channels? - 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 sycl kernel) - colors = torch.cat( - [ - colors[..., :-1], - torch.empty(*colors.shape[:-1], padded_channels, device=device), - colors[..., -1:], - ], - dim=-1, - ) - if backgrounds is not None: - backgrounds = torch.cat( - [ - backgrounds, - torch.zeros( - *backgrounds.shape[:-1], padded_channels, device=device - ), - ], - dim=-1, - ) - else: - padded_channels = 0 - tile_height, tile_width = isect_offsets.shape[-2:] - assert ( - tile_height * tile_size >= image_height - ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" - assert ( - tile_width * tile_size >= image_width - ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - - ( - render_colors, - render_alphas, - render_normals, - render_distort, - render_median, - ) = _RasterizeToPixels2DGS.apply( - means2d.contiguous(), - ray_transforms.contiguous(), - colors.contiguous(), - opacities.contiguous(), - normals.contiguous(), - densify.contiguous(), - backgrounds, - masks, - image_width, - image_height, - tile_size, - isect_offsets.contiguous(), - flatten_ids.contiguous(), - absgrad, - distloss, - ) - - if padded_channels > 0: - render_colors = torch.cat( - [render_colors[..., : -padded_channels - 1], render_colors[..., -1:]], - dim=-1, - ) - - return render_colors, render_alphas, render_normals, render_distort, render_median - - -@torch.no_grad() -def rasterize_to_indices_in_range_2dgs( - range_start: int, - range_end: int, - transmittances: Tensor, # [..., image_height, image_width] - means2d: Tensor, # [..., N, 2] - ray_transforms: Tensor, # [..., N, 3, 3] - opacities: Tensor, # [..., N] - image_width: int, - image_height: int, - tile_size: int, - isect_offsets: Tensor, - flatten_ids: Tensor, -) -> Tuple[Tensor, Tensor, Tensor]: - """Rasterizes a batch of Gaussians to images but only returns the indices. - - .. note:: - - This function supports iterative rasterization, in which each call of this function - will rasterize a batch of Gaussians from near to far, defined by `[range_start, range_end)`. - If a one-step full rasterization is desired, set `range_start` to 0 and `range_end` to a really - large number, e.g, 1e10. - - Args: - range_start: The start batch of Gaussians to be rasterized (inclusive). - range_end: The end batch of Gaussians to be rasterized (exclusive). - transmittances: Currently transmittances. [..., image_height, image_width] - means2d: Projected Gaussian means. [..., N, 2] - ray_transforms: transformation matrices that transforms xy-planes in pixel spaces into splat coordinates. [..., N, 3, 3] - opacities: Gaussian opacities that support per-view values. [..., N] - image_width: Image width. - image_height: Image height. - tile_size: Tile size. - isect_offsets: Intersection offsets outputs from `isect_offset_encode()`. [..., tile_height, tile_width] - flatten_ids: The global flatten indices in [I * N] from `isect_tiles()`. [n_isects] - - Returns: - A tuple: - - - **Gaussian ids**. Gaussian ids for the pixel intersection. A flattened list of shape [M]. - - **Pixel ids**. pixel indices (row-major). A flattened list of shape [M]. - - **Camera ids**. Camera indices. A flattened list of shape [M]. - - **Batch ids**. Batch indices. A flattened list of shape [M]. - """ - - image_dims = means2d.shape[:-2] - tile_height, tile_width = isect_offsets.shape[-2:] - N = means2d.shape[-2] - assert transmittances.shape == image_dims + ( - image_height, - image_width, - ), transmittances.shape - assert means2d.shape == image_dims + (N, 2), means2d.shape - assert ray_transforms.shape == image_dims + (N, 3, 3), ray_transforms.shape - assert opacities.shape == image_dims + (N,), opacities.shape - assert isect_offsets.shape == image_dims + ( - tile_height, - tile_width, - ), isect_offsets.shape - assert ( - tile_height * tile_size >= image_height - ), f"Assert Failed: {tile_height} * {tile_size} >= {image_height}" - assert ( - tile_width * tile_size >= image_width - ), f"Assert Failed: {tile_width} * {tile_size} >= {image_width}" - - out_gauss_ids, out_indices = _make_lazy_sycl_func("rasterize_to_indices_2dgs")( - range_start, - range_end, - transmittances.contiguous(), - means2d.contiguous(), - ray_transforms.contiguous(), - opacities.contiguous(), - image_width, - image_height, - tile_size, - isect_offsets.contiguous(), - flatten_ids.contiguous(), - ) - out_pixel_ids = out_indices % (image_width * image_height) - out_image_ids = out_indices // (image_width * image_height) - return out_gauss_ids, out_pixel_ids, out_image_ids - - -class _RasterizeToPixels2DGS(torch.autograd.Function): - """Rasterize gaussians 2DGS""" - - @staticmethod - def forward( - ctx, - means2d: Tensor, - ray_transforms: Tensor, - colors: Tensor, - opacities: Tensor, - normals: Tensor, - densify: Tensor, - backgrounds: Tensor, - masks: Tensor, - width: int, - height: int, - tile_size: int, - isect_offsets: Tensor, - flatten_ids: Tensor, - absgrad: bool, - distloss: bool, - ) -> Tuple[Tensor, Tensor]: - ( - render_colors, - render_alphas, - render_normals, - render_distort, - render_median, - last_ids, - median_ids, - ) = _make_lazy_sycl_func("rasterize_to_pixels_2dgs_fwd")( - means2d, - ray_transforms, - colors, - opacities, - normals, - backgrounds, - masks, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - ) - - ctx.save_for_backward( - means2d, - ray_transforms, - colors, - opacities, - normals, - densify, - backgrounds, - masks, - isect_offsets, - flatten_ids, - render_colors, - render_alphas, - last_ids, - median_ids, - ) - ctx.width = width - ctx.height = height - ctx.tile_size = tile_size - ctx.absgrad = absgrad - ctx.distloss = distloss - - # double to float - render_alphas = render_alphas.float() - return ( - render_colors, - render_alphas, - render_normals, - render_distort, - render_median, - ) - - @staticmethod - def backward( - ctx, - v_render_colors: Tensor, - v_render_alphas: Tensor, - v_render_normals: Tensor, - v_render_distort: Tensor, - v_render_median: Tensor, - ): - - ( - means2d, - ray_transforms, - colors, - opacities, - normals, - densify, - backgrounds, - masks, - isect_offsets, - flatten_ids, - render_colors, - render_alphas, - last_ids, - median_ids, - ) = ctx.saved_tensors - width = ctx.width - height = ctx.height - tile_size = ctx.tile_size - absgrad = ctx.absgrad - - ( - v_means2d_abs, - v_means2d, - v_ray_transforms, - v_colors, - v_opacities, - v_normals, - v_densify, - ) = _make_lazy_sycl_func("rasterize_to_pixels_2dgs_bwd")( - means2d, - ray_transforms, - colors, - opacities, - normals, - densify, - backgrounds, - masks, - width, - height, - tile_size, - isect_offsets, - flatten_ids, - render_colors, - render_alphas, - last_ids, - median_ids, - v_render_colors.contiguous(), - v_render_alphas.contiguous(), - v_render_normals.contiguous(), - v_render_distort.contiguous(), - v_render_median.contiguous(), - absgrad, - ) - torch.xpu.synchronize() - if absgrad: - means2d.absgrad = v_means2d_abs - - if ctx.needs_input_grad[6]: - v_backgrounds = (v_render_colors * (1.0 - render_alphas).float()).sum( - dim=(-3, -2) - ) - else: - v_backgrounds = None - - return ( - v_means2d, - v_ray_transforms, - v_colors, - v_opacities, - v_normals, - v_densify, - v_backgrounds, - None, - None, - None, - None, - None, - None, - None, - None, - ) diff --git a/setup.py b/setup.py index 7aab77a1..f64b12cf 100644 --- a/setup.py +++ b/setup.py @@ -180,7 +180,7 @@ def get_extensions(): keywords="gaussian, splatting, cuda, sycl", url=URL, download_url=f"{URL}/archive/gsplat-{__version__}.tar.gz", - python_requires=">=3.8", # Updated to match your CMake + python_requires=">=3.8", install_requires=[ "ninja", "numpy", diff --git a/tests/test_2dgs.py b/tests/test_2dgs.py index 8721888d..63dd3553 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -97,7 +97,6 @@ def test_projection_2dgs(test_data, batch_dims: Tuple[int, ...]): # TODO (WZ): is the following true for 2dgs as while? # radii is integer so we allow for 1 unit difference valid = ((radii > 0) & (_radii > 0)).all(dim=-1) - valid_expanded = valid.unsqueeze(-1).unsqueeze(-1) torch.testing.assert_close(radii, _radii, rtol=1e-3, atol=1) torch.testing.assert_close(means2d[valid], _means2d[valid], rtol=1e-4, atol=1e-4) torch.testing.assert_close(depths[valid], _depths[valid], rtol=1e-4, atol=1e-4) @@ -141,7 +140,7 @@ def test_projection_2dgs(test_data, batch_dims: Tuple[int, ...]): 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) From cabd3423786d389e70c37a825c51565da3024578 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Mon, 24 Nov 2025 12:00:01 -0800 Subject: [PATCH 35/56] contiguous() tensors to fused_ssim Disable fused Adam for sycl - requires FP64 (?) simplify backend selection logic. Remove GSPLAT_BACKEND env var. Add device_guard to ensure common device for kernels. Stricter kernel input tensor checks. Update tests, cuda only code. --- examples/requirements.txt | 4 +- examples/simple_trainer.py | 13 ++++-- examples/simple_trainer_2dgs.py | 3 +- examples/utils.py | 2 +- gsplat/__init__.py | 12 +---- gsplat/_helper.py | 4 +- gsplat/rendering.py | 6 +-- gsplat/sycl/include/Common.h | 16 ++++++- gsplat/sycl/src/intersect_offset.cpp | 3 +- gsplat/sycl/src/intersect_tile.cpp | 12 +++-- gsplat/sycl/src/projection_2dgs_fused_bwd.cpp | 23 +++++----- gsplat/sycl/src/projection_2dgs_fused_fwd.cpp | 11 +++-- .../src/projection_ewa_3dgs_fused_bwd.cpp | 27 +++++------ .../src/projection_ewa_3dgs_fused_fwd.cpp | 16 ++++--- .../src/projection_ewa_3dgs_packed_bwd.cpp | 37 +++++++++------ .../src/projection_ewa_3dgs_packed_fwd.cpp | 35 ++++++-------- gsplat/sycl/src/projection_ewa_simple_bwd.cpp | 11 +++-- gsplat/sycl/src/projection_ewa_simple_fwd.cpp | 7 +-- .../src/quat_scale_to_covar_preci_bwd.cpp | 9 ++-- .../src/quat_scale_to_covar_preci_fwd.cpp | 5 +- .../sycl/src/rasterize_to_pixels_2dgs_bwd.cpp | 46 ++++++++++--------- .../sycl/src/rasterize_to_pixels_2dgs_fwd.cpp | 21 +++++---- .../sycl/src/rasterize_to_pixels_3dgs_bwd.cpp | 26 ++++++----- .../sycl/src/rasterize_to_pixels_3dgs_fwd.cpp | 17 +++---- gsplat/sycl/src/relocation.cpp | 6 +++ gsplat/sycl/src/spherical_harmonics_bwd.cpp | 9 ++-- gsplat/sycl/src/spherical_harmonics_fwd.cpp | 15 ++---- gsplat/utils.py | 4 +- profiling/main.py | 13 +++--- tests/_test_distributed.py | 21 +++++---- tests/test_2dgs.py | 5 +- tests/test_basic.py | 5 +- 32 files changed, 242 insertions(+), 202 deletions(-) diff --git a/examples/requirements.txt b/examples/requirements.txt index bbf4fede..741566ef 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/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe +git+https://github.com/rahul-goel/fused-ssim@88169c51c22973ad8fd2429c3298d8356fdd5dc8 +git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe splines diff --git a/examples/simple_trainer.py b/examples/simple_trainer.py index 3b247e13..bc901915 100644 --- a/examples/simple_trainer.py +++ b/examples/simple_trainer.py @@ -21,7 +21,8 @@ generate_interpolated_path, generate_spiral_path, ) -from fusedssim_sycl import fusedssim + +from fused_ssim import fused_ssim from torch import Tensor from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.tensorboard import SummaryWriter @@ -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, torch_acc +from gsplat import export_splats, torch_acc, BACKEND from gsplat.compression import PngCompression from gsplat.distributed import cli from gsplat.optimizers import SelectiveAdam @@ -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 } @@ -682,8 +683,10 @@ def train(self): # loss l1loss = F.l1_loss(colors, pixels) - ssimloss = 1.0 - fusedssim( - colors.permute(0, 3, 1, 2), pixels.permute(0, 3, 1, 2), padding="valid" + ssimloss = 1.0 - fused_ssim( + colors.permute(0, 3, 1, 2).contiguous(), + pixels.permute(0, 3, 1, 2).contiguous(), + padding="valid", ) loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda if cfg.depth_loss: diff --git a/examples/simple_trainer_2dgs.py b/examples/simple_trainer_2dgs.py index 2afda112..c2d533bd 100644 --- a/examples/simple_trainer_2dgs.py +++ b/examples/simple_trainer_2dgs.py @@ -591,7 +591,8 @@ def train(self): # loss l1loss = F.l1_loss(colors, pixels) ssimloss = 1.0 - self.ssim( - pixels.permute(0, 3, 1, 2), colors.permute(0, 3, 1, 2) + pixels.permute(0, 3, 1, 2).contiguous(), + colors.permute(0, 3, 1, 2).contiguous(), ) loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda if cfg.depth_loss: 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/gsplat/__init__.py b/gsplat/__init__.py index 76f98316..161af67e 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -4,20 +4,12 @@ BACKEND: str = "" torch_acc = torch.cpu -_force_backend = os.getenv("GSPLAT_BACKEND", "").lower() -if _force_backend == "cuda" or (_force_backend == "" and torch.cuda.is_available()): +if torch.cuda.is_available(): BACKEND = "cuda" torch_acc = torch.cuda print("gsplat: Using CUDA backend.", file=sys.stderr) - -if ( - not BACKEND - and _force_backend in ("sycl", "xpu") - or _force_backend == "" - and hasattr(torch, "xpu") - and torch.xpu.is_available() -): +elif hasattr(torch, "xpu") and torch.xpu.is_available(): BACKEND = "sycl" torch_acc = torch.xpu print("gsplat: Using SYCL XPU backend.", file=sys.stderr) diff --git a/gsplat/_helper.py b/gsplat/_helper.py index 86bca052..a4d0a0d2 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=torch.accelerator.current_accelerator(), scene_crop: Tuple[float, float, float, float, float, float] = (-2, -2, -2, 2, 2, 2), scene_grid: int = 1, ): @@ -19,6 +19,8 @@ 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.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/rendering.py b/gsplat/rendering.py index cb88b844..7ac7379f 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -72,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 ) @@ -877,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:: @@ -888,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, diff --git a/gsplat/sycl/include/Common.h b/gsplat/sycl/include/Common.h index d92a978d..87b222aa 100644 --- a/gsplat/sycl/include/Common.h +++ b/gsplat/sycl/include/Common.h @@ -10,11 +10,23 @@ 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) + { \ + 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 diff --git a/gsplat/sycl/src/intersect_offset.cpp b/gsplat/sycl/src/intersect_offset.cpp index 3d5dc95b..0f6c7da6 100644 --- a/gsplat/sycl/src/intersect_offset.cpp +++ b/gsplat/sycl/src/intersect_offset.cpp @@ -14,7 +14,8 @@ at::Tensor intersect_offset( const uint32_t tile_width, const uint32_t tile_height ) { - CHECK_CONTIGUOUS(isect_ids); + DEVICE_GUARD(isect_ids); + CHECK_INPUT(isect_ids); const uint32_t C = I; auto options = isect_ids.options().dtype(at::kInt); diff --git a/gsplat/sycl/src/intersect_tile.cpp b/gsplat/sycl/src/intersect_tile.cpp index 5fe318b3..8d9e9f5c 100644 --- a/gsplat/sycl/src/intersect_tile.cpp +++ b/gsplat/sycl/src/intersect_tile.cpp @@ -1,4 +1,5 @@ #include +#include #include @@ -21,13 +22,14 @@ std::tuple intersect_tile( const bool sort, const bool segmented ) { - CHECK_CONTIGUOUS(means2d); - CHECK_CONTIGUOUS(radii); - CHECK_CONTIGUOUS(depths); + DEVICE_GUARD(means2d); + CHECK_INPUT(means2d); + CHECK_INPUT2(radii, means2d); + CHECK_INPUT2(depths, means2d); if (image_ids.has_value()) - CHECK_CONTIGUOUS(image_ids.value()); + CHECK_INPUT2(image_ids.value(), means2d); if (gaussian_ids.has_value()) - CHECK_CONTIGUOUS(gaussian_ids.value()); + CHECK_INPUT2(gaussian_ids.value(), means2d); const bool packed = segmented; const uint32_t C = I; diff --git a/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp index 3756ffc7..796b7834 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_bwd.cpp @@ -26,17 +26,18 @@ projection_2dgs_fused_bwd( const at::Tensor v_ray_transforms, // [..., C, N, 3, 3] const bool viewmats_requires_grad ) { - CHECK_CONTIGUOUS(means); - CHECK_CONTIGUOUS(quats); - CHECK_CONTIGUOUS(scales); - CHECK_CONTIGUOUS(viewmats); - CHECK_CONTIGUOUS(Ks); - CHECK_CONTIGUOUS(radii); - CHECK_CONTIGUOUS(ray_transforms); - CHECK_CONTIGUOUS(v_means2d); - CHECK_CONTIGUOUS(v_depths); - CHECK_CONTIGUOUS(v_normals); - CHECK_CONTIGUOUS(v_ray_transforms); + 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]" diff --git a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp index b39f4e40..c97b8423 100644 --- a/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_2dgs_fused_fwd.cpp @@ -25,11 +25,12 @@ projection_2dgs_fused_fwd( const float far_plane, const float radius_clip ) { - CHECK_CONTIGUOUS(means); - CHECK_CONTIGUOUS(quats); - CHECK_CONTIGUOUS(scales); - CHECK_CONTIGUOUS(viewmats); - CHECK_CONTIGUOUS(Ks); + 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]" diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp index b20314d9..6510bad2 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_bwd.cpp @@ -30,25 +30,26 @@ projection_ewa_3dgs_fused_bwd( const at::optional v_compensations, // [..., C, N] optional const bool viewmats_requires_grad ) { + DEVICE_GUARD(means); // Input validation - CHECK_CONTIGUOUS(means); + CHECK_INPUT(means); if (covars.has_value()) - CHECK_CONTIGUOUS(covars.value()); + CHECK_INPUT2(covars.value(), means); if (quats.has_value()) - CHECK_CONTIGUOUS(quats.value()); + CHECK_INPUT2(quats.value(), means); if (scales.has_value()) - CHECK_CONTIGUOUS(scales.value()); - CHECK_CONTIGUOUS(viewmats); - CHECK_CONTIGUOUS(Ks); - CHECK_CONTIGUOUS(radii); - CHECK_CONTIGUOUS(conics); + 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_CONTIGUOUS(compensations.value()); - CHECK_CONTIGUOUS(v_means2d); - CHECK_CONTIGUOUS(v_depths); - CHECK_CONTIGUOUS(v_conics); + 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_CONTIGUOUS(v_compensations.value()); + CHECK_INPUT2(v_compensations.value(), means); // Dimensions const uint32_t N = means.size(-2); diff --git a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp index fd49d834..9dbeeea9 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_fused_fwd.cpp @@ -29,17 +29,19 @@ projection_ewa_3dgs_fused_fwd( const bool calc_compensations, const CameraModelType camera_model ) { - CHECK_CONTIGUOUS(means); - CHECK_CONTIGUOUS(viewmats); - CHECK_CONTIGUOUS(Ks); + DEVICE_GUARD(means); + // Input validation + CHECK_INPUT(means); + CHECK_INPUT2(viewmats, means); + CHECK_INPUT2(Ks, means); if (covars.has_value()) - CHECK_CONTIGUOUS(covars.value()); + CHECK_INPUT2(covars.value(), means); if (quats.has_value()) - CHECK_CONTIGUOUS(quats.value()); + CHECK_INPUT2(quats.value(), means); if (scales.has_value()) - CHECK_CONTIGUOUS(scales.value()); + CHECK_INPUT2(scales.value(), means); if (opacities.has_value()) - CHECK_CONTIGUOUS(opacities.value()); + CHECK_INPUT2(opacities.value(), means); TORCH_CHECK( means.dim() >= 2, "means must have at least 2 dimensions [..., N, 3]" diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp index 06d9e4f4..45b84d2b 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_bwd.cpp @@ -33,21 +33,28 @@ projection_ewa_3dgs_packed_bwd( const bool viewmats_requires_grad, const bool sparse_grad ) { - TORCH_CHECK( - means.is_contiguous(), "Input 'means' tensor must be contiguous." - ); - TORCH_CHECK( - viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous." - ); - TORCH_CHECK(Ks.is_contiguous(), "Input 'Ks' tensor must be contiguous."); - TORCH_CHECK( - batch_ids.is_contiguous(), - "Input 'batch_ids' tensor must be contiguous." - ); - TORCH_CHECK( - means.device().type() == at::kXPU, - "Input tensors must be on XPU device." - ); + 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); diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp index 339c516f..c070ddc6 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -33,18 +33,20 @@ projection_ewa_3dgs_packed_fwd( const bool calc_compensations, const CameraModelType camera_model ) { - TORCH_CHECK( - means.is_contiguous(), "Input 'means' tensor must be contiguous." - ); - TORCH_CHECK( - viewmats.is_contiguous(), "Input 'viewmats' tensor must be contiguous." - ); - TORCH_CHECK(Ks.is_contiguous(), "Input 'Ks' tensor must be contiguous."); - TORCH_CHECK( - means.device().type() == at::kXPU, - "Input tensors must be on XPU device." - ); - + 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); @@ -83,10 +85,8 @@ projection_ewa_3dgs_packed_fwd( ); } - // --- Start of Correction --- // Allocate block_cnts as kInt, which the kernel expects. at::Tensor block_cnts = at::empty({(long)n_blocks}, int_opts); - // --- End of Correction --- auto &d_queue = at::xpu::getCurrentXPUStream().queue(); sycl::range<2> local_range(1, N_THREADS_PACKED); @@ -125,10 +125,7 @@ projection_ewa_3dgs_packed_fwd( (scalar_t)radius_clip, camera_model, nullptr, // block_accum - // --- Start of Correction --- - // Pass the int32_t pointer directly, no cast needed. block_cnts.data_ptr(), - // --- End of Correction --- nullptr, nullptr, nullptr, @@ -144,11 +141,9 @@ projection_ewa_3dgs_packed_fwd( } ); - // --- Start of Correction --- // 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); - // --- End of Correction --- int64_t nnz = 0; if (n_blocks > 0) { @@ -236,6 +231,7 @@ projection_ewa_3dgs_packed_fwd( } return std::make_tuple( + indptr, batch_ids, camera_ids, gaussian_ids, @@ -243,7 +239,6 @@ projection_ewa_3dgs_packed_fwd( means2d, depths, conics, - indptr, compensations ); } diff --git a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp index a19bc7fc..abd70cf7 100644 --- a/gsplat/sycl/src/projection_ewa_simple_bwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_bwd.cpp @@ -16,11 +16,12 @@ std::tuple projection_ewa_simple_bwd( const at::Tensor v_means2d, // [..., C, N, 2] const at::Tensor v_covars2d // [..., C, N, 2, 2] ) { - CHECK_CONTIGUOUS(means); - CHECK_CONTIGUOUS(covars); - CHECK_CONTIGUOUS(Ks); - CHECK_CONTIGUOUS(v_means2d); - CHECK_CONTIGUOUS(v_covars2d); + 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); diff --git a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp index 2a692816..604af3f2 100644 --- a/gsplat/sycl/src/projection_ewa_simple_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_simple_fwd.cpp @@ -14,9 +14,10 @@ std::tuple projection_ewa_simple_fwd( const uint32_t height, const CameraModelType camera_model ) { - CHECK_CONTIGUOUS(means); - CHECK_CONTIGUOUS(covars); - CHECK_CONTIGUOUS(Ks); + 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]" ); diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp index 9a3d1eba..a774764f 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_bwd.cpp @@ -12,13 +12,14 @@ std::tuple quat_scale_to_covar_preci_bwd( const at::optional v_covars, // [..., 3, 3] or [..., 6] const at::optional v_precis // [..., 3, 3] or [..., 6] ) { - CHECK_CONTIGUOUS(quats); - CHECK_CONTIGUOUS(scales); + DEVICE_GUARD(quats); + CHECK_INPUT(quats); + CHECK_INPUT2(scales, quats); if (v_covars.has_value()) { - CHECK_CONTIGUOUS(v_covars.value()); + CHECK_INPUT2(v_covars.value(), quats); } if (v_precis.has_value()) { - CHECK_CONTIGUOUS(v_precis.value()); + CHECK_INPUT2(v_precis.value(), quats); } TORCH_CHECK( v_covars.has_value() || v_precis.has_value(), diff --git a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp index d20be021..c08640ae 100644 --- a/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp +++ b/gsplat/sycl/src/quat_scale_to_covar_preci_fwd.cpp @@ -12,8 +12,9 @@ std::tuple quat_scale_to_covar_preci_fwd( const bool compute_preci, const bool triu ) { - CHECK_CONTIGUOUS(quats); - CHECK_CONTIGUOUS(scales); + 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" diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp index 2bc20ce0..9858d9e4 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_bwd.cpp @@ -186,28 +186,30 @@ rasterize_to_pixels_2dgs_bwd( const at::Tensor v_render_median, // [..., image_height, image_width] bool absgrad ) { - // Check input tensors are contiguous - CHECK_CONTIGUOUS(means2d); - CHECK_CONTIGUOUS(ray_transforms); - CHECK_CONTIGUOUS(colors); - CHECK_CONTIGUOUS(opacities); - CHECK_CONTIGUOUS(normals); - CHECK_CONTIGUOUS(densify); - CHECK_CONTIGUOUS(tile_offsets); - CHECK_CONTIGUOUS(flatten_ids); - CHECK_CONTIGUOUS(render_colors); - CHECK_CONTIGUOUS(render_alphas); - CHECK_CONTIGUOUS(last_ids); - CHECK_CONTIGUOUS(median_ids); - CHECK_CONTIGUOUS(v_render_colors); - CHECK_CONTIGUOUS(v_render_alphas); - CHECK_CONTIGUOUS(v_render_normals); - CHECK_CONTIGUOUS(v_render_distort); - CHECK_CONTIGUOUS(v_render_median); - if (backgrounds.has_value()) - CHECK_CONTIGUOUS(backgrounds.value()); - if (masks.has_value()) - CHECK_CONTIGUOUS(masks.value()); + 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); diff --git a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp index 960ffe11..0cfa20e3 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_2dgs_fwd.cpp @@ -134,18 +134,19 @@ rasterize_to_pixels_2dgs_fwd( const at::Tensor tile_offsets, // [..., tile_height, tile_width] const at::Tensor flatten_ids // [n_isects] ) { - // Check input tensors are contiguous - CHECK_CONTIGUOUS(means2d); - CHECK_CONTIGUOUS(ray_transforms); - CHECK_CONTIGUOUS(colors); - CHECK_CONTIGUOUS(opacities); - CHECK_CONTIGUOUS(normals); - CHECK_CONTIGUOUS(tile_offsets); - CHECK_CONTIGUOUS(flatten_ids); + 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_CONTIGUOUS(backgrounds.value()); + CHECK_INPUT2(backgrounds.value(), means2d); if (masks.has_value()) - CHECK_CONTIGUOUS(masks.value()); + CHECK_INPUT2(masks.value(), means2d); // Get dimensions bool packed = means2d.dim() == 2; diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp index 3fc64713..629bafa9 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_bwd.cpp @@ -153,20 +153,22 @@ rasterize_to_pixels_3dgs_bwd( // options bool absgrad ) { - CHECK_CONTIGUOUS(means2d); - CHECK_CONTIGUOUS(conics); - CHECK_CONTIGUOUS(colors); - CHECK_CONTIGUOUS(opacities); - CHECK_CONTIGUOUS(tile_offsets); - CHECK_CONTIGUOUS(flatten_ids); - CHECK_CONTIGUOUS(render_alphas); - CHECK_CONTIGUOUS(last_ids); - CHECK_CONTIGUOUS(v_render_colors); - CHECK_CONTIGUOUS(v_render_alphas); + 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_CONTIGUOUS(backgrounds.value()); + CHECK_INPUT2(backgrounds.value(), means2d); if (masks.has_value()) - CHECK_CONTIGUOUS(masks.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"); diff --git a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp index 2875f900..802adfc5 100644 --- a/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp +++ b/gsplat/sycl/src/rasterize_to_pixels_3dgs_fwd.cpp @@ -115,16 +115,17 @@ std::tuple rasterize_to_pixels_3dgs_fwd( const at::Tensor tile_offsets, const at::Tensor flatten_ids ) { - CHECK_CONTIGUOUS(means2d); - CHECK_CONTIGUOUS(conics); - CHECK_CONTIGUOUS(colors); - CHECK_CONTIGUOUS(opacities); - CHECK_CONTIGUOUS(tile_offsets); - CHECK_CONTIGUOUS(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_CONTIGUOUS(backgrounds.value()); + CHECK_INPUT2(backgrounds.value(), means2d); if (masks.has_value()) - CHECK_CONTIGUOUS(masks.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"); diff --git a/gsplat/sycl/src/relocation.cpp b/gsplat/sycl/src/relocation.cpp index fcb31d8d..17777a7d 100644 --- a/gsplat/sycl/src/relocation.cpp +++ b/gsplat/sycl/src/relocation.cpp @@ -15,6 +15,12 @@ std::tuple relocation( 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) diff --git a/gsplat/sycl/src/spherical_harmonics_bwd.cpp b/gsplat/sycl/src/spherical_harmonics_bwd.cpp index ccba983d..acb8fd1f 100644 --- a/gsplat/sycl/src/spherical_harmonics_bwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_bwd.cpp @@ -14,11 +14,12 @@ std::tuple spherical_harmonics_bwd( const at::Tensor v_colors, // [..., 3] bool compute_v_dirs ) { - CHECK_CONTIGUOUS(dirs); - CHECK_CONTIGUOUS(coeffs); - CHECK_CONTIGUOUS(v_colors); + DEVICE_GUARD(dirs); + CHECK_INPUT(dirs); + CHECK_INPUT2(coeffs, dirs); + CHECK_INPUT2(v_colors, dirs); if (masks.has_value()) { - CHECK_CONTIGUOUS(masks.value()); + CHECK_INPUT2(masks.value(), dirs); } TORCH_CHECK(v_colors.size(-1) == 3, "v_colors must have last dimension 3"); diff --git a/gsplat/sycl/src/spherical_harmonics_fwd.cpp b/gsplat/sycl/src/spherical_harmonics_fwd.cpp index 53f11cff..3885af3f 100644 --- a/gsplat/sycl/src/spherical_harmonics_fwd.cpp +++ b/gsplat/sycl/src/spherical_harmonics_fwd.cpp @@ -11,19 +11,12 @@ at::Tensor spherical_harmonics_fwd( const at::Tensor coeffs, // [..., K, 3] const at::optional masks // [...] ) { - TORCH_CHECK( - dirs.is_contiguous(), "Input 'dirs' tensor must be contiguous." - ); - TORCH_CHECK( - coeffs.is_contiguous(), "Input 'coeffs' tensor must be contiguous." - ); + DEVICE_GUARD(dirs); + CHECK_INPUT(dirs); + CHECK_INPUT2(coeffs, dirs); if (masks.has_value()) { - TORCH_CHECK( - masks.value().is_contiguous(), - "Input 'masks' tensor must be contiguous." - ); + CHECK_INPUT2(masks.value(), dirs); } - TORCH_CHECK( dirs.size(-1) == 3, "Input 'dirs' tensor must have the last dimension of size 3." diff --git a/gsplat/utils.py b/gsplat/utils.py index 4924091f..61c60eae 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=torch.accelerator.current_accelerator()) -> Tensor: """Create OpenGL-style projection matrix""" tanHalfFovY = math.tan((fovY / 2)) tanHalfFovX = math.tan((fovX / 2)) @@ -241,6 +241,8 @@ def get_projection_matrix(znear, zfar, fovX, fovY, device="cuda"): right = tanHalfFovX * znear left = -right + if device is None: + device = torch.device("cpu") P = torch.zeros(4, 4, device=device) z_sign = 1.0 diff --git a/profiling/main.py b/profiling/main.py index 028295f0..e4e636c1 100644 --- a/profiling/main.py +++ b/profiling/main.py @@ -6,12 +6,13 @@ ``` """ +import os import time import torch from typing_extensions import Callable, Literal -from gsplat import torch_acc, BACKEND +from gsplat import __version__, torch_acc, BACKEND from gsplat._helper import load_test_data from gsplat.distributed import cli from gsplat.rendering import rasterization @@ -51,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, @@ -61,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) @@ -181,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 @@ -212,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 @@ -243,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 diff --git a/tests/_test_distributed.py b/tests/_test_distributed.py index ae03f9c9..d1177198 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,13 @@ 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 +29,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 +48,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 +72,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 +107,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 63dd3553..9475b617 100644 --- a/tests/test_2dgs.py +++ b/tests/test_2dgs.py @@ -16,6 +16,9 @@ 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, ...]): @@ -134,7 +137,7 @@ def test_projection_2dgs(test_data, batch_dims: Tuple[int, ...]): torch.testing.assert_close(v_means, _v_means, rtol=1e-2, atol=6e-2) -@requires_backend +@requires_cuda @pytest.mark.parametrize("sparse_grad", [False]) @pytest.mark.parametrize("batch_dims", [(), (2,), (1, 2)]) def test_fully_fused_projection_packed_2dgs( diff --git a/tests/test_basic.py b/tests/test_basic.py index 7f7c1a33..0cba2926 100644 --- a/tests/test_basic.py +++ b/tests/test_basic.py @@ -3,8 +3,7 @@ Usage: pytest -s - -# To force a specific backend for testing: +``` """ import math @@ -357,7 +356,7 @@ def test_fully_fused_projection_packed( depths, conics, compensations, - ) = fully_fused_projection( + ) = gsplat.fully_fused_projection( means, covars, None, From 31c484e871396c9ae4c7f08ac200c3d855d22613 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Mon, 24 Nov 2025 12:14:03 -0800 Subject: [PATCH 36/56] Updated fused_ssim to latest including sycl bugfix. --- examples/requirements.txt | 2 +- examples/simple_trainer.py | 6 ++---- examples/simple_trainer_2dgs.py | 3 +-- 3 files changed, 4 insertions(+), 7 deletions(-) diff --git a/examples/requirements.txt b/examples/requirements.txt index 741566ef..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@88169c51c22973ad8fd2429c3298d8356fdd5dc8 +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 bc901915..0f46d2b2 100644 --- a/examples/simple_trainer.py +++ b/examples/simple_trainer.py @@ -295,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=(None if BACKEND == 'sycl' else True), + fused=(None if BACKEND == "sycl" else True), ) for name, _, lr in params } @@ -684,9 +684,7 @@ def train(self): # loss l1loss = F.l1_loss(colors, pixels) ssimloss = 1.0 - fused_ssim( - colors.permute(0, 3, 1, 2).contiguous(), - pixels.permute(0, 3, 1, 2).contiguous(), - padding="valid", + colors.permute(0, 3, 1, 2), pixels.permute(0, 3, 1, 2), padding="valid" ) loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda if cfg.depth_loss: diff --git a/examples/simple_trainer_2dgs.py b/examples/simple_trainer_2dgs.py index c2d533bd..2afda112 100644 --- a/examples/simple_trainer_2dgs.py +++ b/examples/simple_trainer_2dgs.py @@ -591,8 +591,7 @@ def train(self): # loss l1loss = F.l1_loss(colors, pixels) ssimloss = 1.0 - self.ssim( - pixels.permute(0, 3, 1, 2).contiguous(), - colors.permute(0, 3, 1, 2).contiguous(), + pixels.permute(0, 3, 1, 2), colors.permute(0, 3, 1, 2) ) loss = l1loss * (1.0 - cfg.ssim_lambda) + ssimloss * cfg.ssim_lambda if cfg.depth_loss: From a70b2f57f8c8dd255b292b35ed4e139240b33143 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Tue, 25 Nov 2025 10:24:46 -0800 Subject: [PATCH 37/56] empty xpu docs --- docs/Intel XPU.md | 134 +++++++++++++++++++++++++++++++++++++ gsplat/_helper.py | 8 ++- gsplat/relocation.py | 1 + gsplat/utils.py | 8 ++- tests/_test_distributed.py | 1 + 5 files changed, 148 insertions(+), 4 deletions(-) create mode 100644 docs/Intel XPU.md diff --git a/docs/Intel XPU.md b/docs/Intel XPU.md new file mode 100644 index 00000000..e1fb3dbd --- /dev/null +++ b/docs/Intel XPU.md @@ -0,0 +1,134 @@ +# GSPLAT on Intel GPUs + +``gsplat`` supports creation and rendering on Intel GPUs through SYCL kernel backend. This provides support for both integrated (Alder Lake Arc and onward) as well as 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 + python -m pip install torch --index-url https://download.pytorch.org/whl/xpu + ``` + +- **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 tolkit version must match the version used to build PyTorch XPU. Check the PyTorch XPU OneAPI version with: + + pip show intel-cmplr-lib-ur # dependency of torch-xpu + ... + Version: 2025.0.5 + ... + +- Configure your build environment: + + In Linux: + + ```bash + source /opt/intel/oneapi/setvars.sh + ``` + + Or in Windows: + + ```ps1 + cmd /k "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" + powershell + cmd /k "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" + powershell + ``` + +- Finally, build and install the project's Python extension. This step might take some time. + + ```bash + pip install --no-build-isolation . + ``` + + Alternately, you can build a wheel for distribution with: + + ```bash + 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 GPU. + +### 3DGS Reproduced metrics + +| PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + + +| SSIM | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + + +| LPIPS | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + +| Num GSs | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + +### 3DGS Training time and memory + +| Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------------------|---------|--------|---------|--------|---------|--------|--------| +| 7k steps Mem (GB) | | | | | | | | +| 30k steps Mem (GB) | | | | | | | | +| 7k steps time (s) | | | | | | | | +| 30k steps time (s) | | | | | | | | + +### 2DGS Reproduced metrics + +| PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + + +| SSIM | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + + +| LPIPS | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + +| Num GSs | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------|---------|--------|---------|--------|---------|-------|-------| +| 7k steps | | | | | | | | +| 30k steps | | | | | | | | + +### 2DGS Training time and memory + +| Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | +|-----------------------|---------|--------|---------|--------|---------|--------|--------| +| 7k steps Mem (GB) | | | | | | | | +| 30k steps Mem (GB) | | | | | | | | +| 7k steps time (s) | | | | | | | | +| 30k steps time (s) | | | | | | | | \ No newline at end of file diff --git a/gsplat/_helper.py b/gsplat/_helper.py index a4d0a0d2..c02818b8 100644 --- a/gsplat/_helper.py +++ b/gsplat/_helper.py @@ -8,7 +8,7 @@ def load_test_data( data_path: Optional[str] = None, - device=torch.accelerator.current_accelerator(), + device=None, scene_crop: Tuple[float, float, float, float, float, float] = (-2, -2, -2, 2, 2, 2), scene_grid: int = 1, ): @@ -20,7 +20,11 @@ def load_test_data( data = np.load(data_path) height, width = data["height"].item(), data["width"].item() if device is None: - device = torch.device("cpu") + 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/relocation.py b/gsplat/relocation.py index 0f6538b7..0a6600d5 100644 --- a/gsplat/relocation.py +++ b/gsplat/relocation.py @@ -7,6 +7,7 @@ from . import BACKEND from ._wrapper import _make_lazy_device_func + def compute_relocation( opacities: Tensor, # [N] scales: Tensor, # [N, 3] diff --git a/gsplat/utils.py b/gsplat/utils.py index 61c60eae..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=torch.accelerator.current_accelerator()) -> Tensor: +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)) @@ -242,7 +242,11 @@ def get_projection_matrix(znear, zfar, fovX, fovY, device=torch.accelerator.curr left = -right if device is None: - device = torch.device("cpu") + 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/tests/_test_distributed.py b/tests/_test_distributed.py index d1177198..46de2b89 100644 --- a/tests/_test_distributed.py +++ b/tests/_test_distributed.py @@ -15,6 +15,7 @@ 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(local_rank) From bc600981a0d267867e87c26203591985844a0947 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Mon, 26 Jan 2026 12:39:26 -0800 Subject: [PATCH 38/56] Evaluation results on B580 --- docs/Intel XPU.md | 52 +++++++++++++++++++++++------------------------ 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/docs/Intel XPU.md b/docs/Intel XPU.md index e1fb3dbd..7b842378 100644 --- a/docs/Intel XPU.md +++ b/docs/Intel XPU.md @@ -22,7 +22,7 @@ The kernels are optimized and use mixed precision (some data is represented as h - **PyTorch XPU:** Install the PyTorch XPU version. ```bash - python -m pip install torch --index-url https://download.pytorch.org/whl/xpu + python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/xpu ``` - **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. @@ -51,7 +51,7 @@ The kernels are optimized and use mixed precision (some data is represented as h powershell ``` -- Finally, build and install the project's Python extension. This step might take some time. +- Finally, build and install the project's Python extension. ```bash pip install --no-build-isolation . @@ -71,64 +71,64 @@ We evaluate gsplat-xpu on the Mip-NeRF 360 dataset and measure PSNR, SSIM, LPIPS | PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------|---------|--------|---------|--------|---------|-------|-------| -| 7k steps | | | | | | | | -| 30k steps | | | | | | | | +| 7k steps | 24.01 | 29.66 | 27.26 | 26.59 | 28.65 | 28.70 | 26.03 | +| 30k steps | | 31.89 | 29.14 | | 30.90 | 31.06 | | | SSIM | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------|---------|--------|---------|--------|---------|-------|-------| -| 7k steps | | | | | | | | -| 30k steps | | | | | | | | +| 7k steps | 0.6808 | 0.9262 | 0.8865 | 0.8370 | 0.9047 | 0.8945| 0.7378| +| 30k steps | | 0.9446 | 0.9158 | | 0.9318 | 0.9239| | | LPIPS | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------|---------|--------|---------|--------|---------|-------|-------| -| 7k steps | | | | | | | | -| 30k steps | | | | | | | | +| 7k steps | 0.2997 | 0.1462 | 0.1929 | 0.1195 | 0.1220 | 0.2136| 0.2339| +| 30k steps | | 0.1179 | 0.1414 | | 0.08607 | 0.1520| | | Num GSs | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------|---------|--------|---------|--------|---------|-------|-------| -| 7k steps | | | | | | | | -| 30k steps | | | | | | | | +| 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.28 M | 1.27 M | | 1.90 M | 1.63 M| | ### 3DGS Training time and memory | Mip-NeRF 360 scene | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------------------|---------|--------|---------|--------|---------|--------|--------| -| 7k steps Mem (GB) | | | | | | | | -| 30k steps Mem (GB) | | | | | | | | -| 7k steps time (s) | | | | | | | | -| 30k steps time (s) | | | | | | | | +| 7k steps Mem (GB) | 5.846 | 2.009 | 1.730 | 6.177 | 2.737 | 1.861 | 5.921 | +| 30k steps Mem (GB) | | 2.043 | 1.986 | | 2.923 | 2.463 | | +| 7k steps time (s) | 588.5 | 534.4 | 625.3 | 793.6 | 919.1 | 645.1 | 519.3 | +| 30k steps time (s) | | 2612 | 3520 | | 5027 | 3317 | | ### 2DGS Reproduced metrics | PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------|---------|--------|---------|--------|---------|-------|-------| -| 7k steps | | | | | | | | -| 30k steps | | | | | | | | +| 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 | | | | | | | | -| 30k steps | | | | | | | | +| 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 | | | | | | | | -| 30k steps | | | | | | | | +| 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 | | | | | | | | -| 30k steps | | | | | | | | +| 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) | | | | | | | | -| 30k steps Mem (GB) | | | | | | | | -| 7k steps time (s) | | | | | | | | -| 30k steps time (s) | | | | | | | | \ No newline at end of file +| 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 | \ No newline at end of file From 4f022e4f18308c9164681b922b8b460a86be02b3 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Tue, 27 Jan 2026 14:18:13 -0800 Subject: [PATCH 39/56] Windows MSVC fixes --- gsplat/sycl/CMakeLists.txt | 2 +- gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp | 2 +- gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt index 28007357..12f2fc20 100644 --- a/gsplat/sycl/CMakeLists.txt +++ b/gsplat/sycl/CMakeLists.txt @@ -39,7 +39,7 @@ else() "Please ensure PyTorch is installed or set CMAKE_PREFIX_PATH/Torch_DIR manually.") endif() -string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_SHARED_LIBRARY_SUFFIX}) +string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_LINK_LIBRARY_SUFFIX}) execute_process( COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], '${TORCH_PYTHON_LIB_NAME}'))" OUTPUT_STRIP_TRAILING_WHITESPACE diff --git a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp index 8ce95ca5..2d218747 100644 --- a/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp +++ b/gsplat/sycl/include/kernels/PackedProjectionFwdKernel.hpp @@ -86,7 +86,7 @@ template struct PackedProjectionFwdKernel { void operator()(sycl::nd_item<2> item) const { auto group = item.get_group(); - sycl::id<2> group_id = item.get_group_id(); + 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(); diff --git a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp index c070ddc6..168d8d46 100644 --- a/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp +++ b/gsplat/sycl/src/projection_ewa_3dgs_packed_fwd.cpp @@ -227,7 +227,7 @@ projection_ewa_3dgs_packed_fwd( // Set the last element of indptr if (nrows > 0) { - indptr.index_put_({(long)nrows}, nnz); + indptr.index_put_({at::indexing::TensorIndex((int64_t)nrows)}, nnz); } return std::make_tuple( From a9f15365ea331efd58e337377dc268a76262e3c0 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Wed, 4 Feb 2026 20:46:28 +0100 Subject: [PATCH 40/56] Do not install fused-bilagrid --- examples/benchmarks/basic.sh | 6 +++--- examples/benchmarks/basic_2dgs.sh | 6 +++--- examples/requirements.txt | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) 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/requirements.txt b/examples/requirements.txt index a55535fb..73b3951e 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -20,5 +20,5 @@ tensorly pyyaml matplotlib git+https://github.com/rahul-goel/fused-ssim@b42e988db507702aa1198920ec7b36a5aa3f72b2 -git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe +#git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe splines From b315442feb102512d26ec34d7a751791072f288d Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Wed, 4 Feb 2026 11:52:16 -0800 Subject: [PATCH 41/56] Update instructions. --- docs/{Intel XPU.md => Intel_XPU.md} | 84 +++++++++++++++++++++++------ 1 file changed, 68 insertions(+), 16 deletions(-) rename docs/{Intel XPU.md => Intel_XPU.md} (59%) diff --git a/docs/Intel XPU.md b/docs/Intel_XPU.md similarity index 59% rename from docs/Intel XPU.md rename to docs/Intel_XPU.md index 7b842378..2a18b697 100644 --- a/docs/Intel XPU.md +++ b/docs/Intel_XPU.md @@ -22,7 +22,7 @@ The kernels are optimized and use mixed precision (some data is represented as h - **PyTorch XPU:** Install the PyTorch XPU version. ```bash - python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/xpu + pip install torch torchvision --index-url https://download.pytorch.org/whl/xpu ``` - **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. @@ -31,7 +31,7 @@ The kernels are optimized and use mixed precision (some data is represented as h pip show intel-cmplr-lib-ur # dependency of torch-xpu ... - Version: 2025.0.5 + Version: 2025.3.1 ... - Configure your build environment: @@ -42,13 +42,14 @@ The kernels are optimized and use mixed precision (some data is represented as h source /opt/intel/oneapi/setvars.sh ``` - Or in Windows: + Or in Windows, setup your Visual Studio build environment and then OneAPI build environment. For example: ```ps1 cmd /k "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" powershell 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. @@ -60,47 +61,62 @@ The kernels are optimized and use mixed precision (some data is represented as h Alternately, you can build a wheel for distribution with: ```bash - python -m build --no-isolation --wheel + 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 GPU. +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: -### 3DGS Reproduced metrics + ```bash + cd examples + python datasets/download_dataset.py + pip install --no-build-isolation -r requirements.txt + ``` + +The last command will also build and install the `fused-ssim` package. This needs the `--no-build-isolation` option. 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 + 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 | | 31.89 | 29.14 | | 30.90 | 31.06 | | +| 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 | | 0.9446 | 0.9158 | | 0.9318 | 0.9239| | +| 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 | | 0.1179 | 0.1414 | | 0.08607 | 0.1520| | +| 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.28 M | 1.27 M | | 1.90 M | 1.63 M| | +| 30k steps | [^1] | 1.28 M | 1.27 M | [^1] | 1.90 M | 1.63 M| [^1] | -### 3DGS Training time and memory +#### 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) | | 2.043 | 1.986 | | 2.923 | 2.463 | | +| 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) | | 2612 | 3520 | | 5027 | 3317 | | +| 30k steps time (s) | [^1] | 2612 | 3520 | [^1] | 5027 | 3317 | [^1] | -### 2DGS Reproduced metrics +#### 2DGS Reproduced metrics | PSNR | Bicycle | Bonsai | Counter | Garden | Kitchen | Room | Stump | |-----------|---------|--------|---------|--------|---------|-------|-------| @@ -124,11 +140,47 @@ We evaluate gsplat-xpu on the Mip-NeRF 360 dataset and measure PSNR, SSIM, LPIPS | 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 +#### 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 | \ No newline at end of file +| 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 From fa4cdd97dc13e42dc27742434d98595887138bde Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Sat, 7 Feb 2026 09:08:39 +0100 Subject: [PATCH 42/56] Simplify build process. Add pyproject.toml for build dependencies. Simplify CMakeLists.txt TODO: Check Windows. --- MANIFEST.in | 7 +++- docs/Intel_XPU.md | 24 ++++++------- gsplat/__init__.py | 1 - gsplat/rendering.py | 6 ++-- gsplat/sycl/CMakeLists.txt | 69 ++++++++++++++++++-------------------- gsplat/sycl/_backend.py | 2 ++ pyproject.toml | 3 ++ setup.py | 23 +++++++------ 8 files changed, 72 insertions(+), 63 deletions(-) create mode 100644 pyproject.toml 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/docs/Intel_XPU.md b/docs/Intel_XPU.md index 2a18b697..7fd7f09e 100644 --- a/docs/Intel_XPU.md +++ b/docs/Intel_XPU.md @@ -1,6 +1,6 @@ # GSPLAT on Intel GPUs -``gsplat`` supports creation and rendering on Intel GPUs through SYCL kernel backend. This provides support for both integrated (Alder Lake Arc and onward) as well as discrete GPUs (Arc Alchemist and newer, such as the A770 and B580). +`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: @@ -27,12 +27,14 @@ The kernels are optimized and use mixed precision (some data is represented as h - **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 tolkit version must match the version used to build PyTorch XPU. Check the PyTorch XPU OneAPI version with: + **Note:** The OneAPI toolkit version must match the version used to build PyTorch XPU. Check the PyTorch XPU OneAPI version with: - pip show intel-cmplr-lib-ur # dependency of torch-xpu - ... - Version: 2025.3.1 - ... + ```bash + pip show intel-cmplr-lib-ur # dependency of torch-xpu + # ... + # Version: 2025.3.1 + # ... + ``` - Configure your build environment: @@ -45,8 +47,6 @@ The kernels are optimized and use mixed precision (some data is represented as h Or in Windows, setup your Visual Studio build environment and then OneAPI build environment. For example: ```ps1 - cmd /k "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat" - powershell cmd /k "C:\Program Files (x86)\Intel\oneAPI\setvars.bat" powershell $env:DISTUTILS_USE_SDK=1 @@ -55,13 +55,13 @@ The kernels are optimized and use mixed precision (some data is represented as h - Finally, build and install the project's Python extension. ```bash - pip install --no-build-isolation . + pip install --extra-index-url=https://download.pytorch.org/whl/xpu . ``` Alternately, you can build a wheel for distribution with: ```bash - python -m build --no-isolation --wheel . + PIP_EXTRA_INDEX_URL=https://download.pytorch.org/whl/xpu python -m build --no-isolation --wheel . ``` ## Evaluation @@ -71,10 +71,10 @@ We evaluate gsplat-xpu on the Mip-NeRF 360 dataset and measure PSNR, SSIM, LPIPS ```bash cd examples python datasets/download_dataset.py - pip install --no-build-isolation -r requirements.txt + pip install --extra-index-url=https://download.pytorch.org/whl/xpu -r requirements.txt ``` -The last command will also build and install the `fused-ssim` package. This needs the `--no-build-isolation` option. 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: +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 bash benchmarks/basic.sh diff --git a/gsplat/__init__.py b/gsplat/__init__.py index 161af67e..febfbe3d 100644 --- a/gsplat/__init__.py +++ b/gsplat/__init__.py @@ -43,7 +43,6 @@ from .strategy import DefaultStrategy, MCMCStrategy, Strategy from .version import __version__ - __all__ = [ "BACKEND", "torch_acc", diff --git a/gsplat/rendering.py b/gsplat/rendering.py index 7ac7379f..875cecc5 100644 --- a/gsplat/rendering.py +++ b/gsplat/rendering.py @@ -622,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, @@ -650,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, @@ -674,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), diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt index 12f2fc20..b6a03e6b 100644 --- a/gsplat/sycl/CMakeLists.txt +++ b/gsplat/sycl/CMakeLists.txt @@ -16,59 +16,56 @@ endif() find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED) - execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import torch.utils; print(torch.utils.cmake_prefix_path)" + 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) +base_path = os.path.join(lib_paths[0], '${CMAKE_SHARED_LIBRARY_PREFIX}torch_python') +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 - OUTPUT_VARIABLE Torch_DIR_From_Python - RESULT_VARIABLE _torch_path_result - ERROR_QUIET ) -if(NOT _torch_path_result EQUAL 0) - message(WARNING "Failed to get Torch CMake path from Python. " - "Make sure PyTorch is installed in the Python environment: ${Python_EXECUTABLE}") - set(Torch_DIR_From_Python "") + +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. " - "Please ensure PyTorch is installed or set CMAKE_PREFIX_PATH/Torch_DIR manually.") + message(FATAL_ERROR "Could not find Torch via Python introspection.") endif() -string(CONCAT TORCH_PYTHON_LIB_NAME ${CMAKE_SHARED_LIBRARY_PREFIX} "torch_python" ${CMAKE_LINK_LIBRARY_SUFFIX}) -execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import os; from torch.utils import cpp_extension; print(os.path.join(cpp_extension.library_paths(True)[0], '${TORCH_PYTHON_LIB_NAME}'))" - OUTPUT_STRIP_TRAILING_WHITESPACE - OUTPUT_VARIABLE TORCH_PYTHON_LIB -) - -# Add a fallback to find the file if it exists with a different extension -if (NOT EXISTS "${TORCH_PYTHON_LIB}") - execute_process( - COMMAND "${Python_EXECUTABLE}" -c "import os, glob; from torch.utils import cpp_extension; base_path = os.path.join(cpp_extension.library_paths(True)[0], '${CMAKE_SHARED_LIBRARY_PREFIX}torch_python'); matches = glob.glob(base_path + '.*'); print(matches[0] if matches else '')" - OUTPUT_STRIP_TRAILING_WHITESPACE - OUTPUT_VARIABLE TORCH_PYTHON_LIB_FALLBACK - ) - if (EXISTS "${TORCH_PYTHON_LIB_FALLBACK}") - set(TORCH_PYTHON_LIB "${TORCH_PYTHON_LIB_FALLBACK}") - message(STATUS "Found torch_python library using fallback: ${TORCH_PYTHON_LIB}") - endif() -endif() - - if (NOT EXISTS "${TORCH_PYTHON_LIB}") - message(FATAL_ERROR "Could not find ${TORCH_PYTHON_LIB_NAME} at ${TORCH_PYTHON_LIB}. Please check your PyTorch installation.") + message(FATAL_ERROR "Could not find torch_python library.") else() message(STATUS "Found torch_python library at: ${TORCH_PYTHON_LIB}") endif() - -set(PYBIND11_FINDPYTHON ON) -find_package(pybind11 CONFIG REQUIRED HINTS "${Python_SITELIB}/pybind11/share/cmake/pybind11") +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 diff --git a/gsplat/sycl/_backend.py b/gsplat/sycl/_backend.py index fbc8e5f0..abe4e60e 100644 --- a/gsplat/sycl/_backend.py +++ b/gsplat/sycl/_backend.py @@ -22,3 +22,5 @@ if os.name == "nt": for dp in dllpath: dp.close() + +__all__ = ["_C"] 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 f64b12cf..39269102 100644 --- a/setup.py +++ b/setup.py @@ -51,22 +51,26 @@ class SyclBuildExtension(BuildExtension): def run(self): print("--- Running SYCL build via CMake ---") + import shutil + 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) + cmake_args = [ + "cmake", + "-G", + "Ninja", + 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", - "-G", - "Ninja", - "-DCMAKE_BUILD_TYPE=Release", - f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={os.path.join(install_dir, 'gsplat')}", - sycl_dir, - ], + cmake_args, cwd=build_dir, ) sp.check_call( @@ -201,7 +205,6 @@ def get_extensions(): "build", "twine", ], - "sycl": ["pybind11>=2.10"], }, ext_modules=ext_modules, cmdclass=cmdclass, From 3cb826210a4f1fa5af9d4741493a033cc460ccc4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 23 Feb 2026 18:06:22 +0000 Subject: [PATCH 43/56] Initial plan From 02d094ce697ebdefb6e668363111c170fc256d2d Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 23 Feb 2026 18:20:49 +0000 Subject: [PATCH 44/56] Add SYCL/XPU wheel build and publish GitHub Actions workflows Co-authored-by: ssheorey <41028320+ssheorey@users.noreply.github.com> --- .github/workflows/building_xpu.yml | 86 +++++++++++++++++++++++++++++ .github/workflows/publish_xpu.yml | 88 ++++++++++++++++++++++++++++++ .github/workflows/xpu/Linux.sh | 14 +++++ setup.py | 2 +- 4 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 .github/workflows/building_xpu.yml create mode 100644 .github/workflows/publish_xpu.yml create mode 100644 .github/workflows/xpu/Linux.sh diff --git a/.github/workflows/building_xpu.yml b/.github/workflows/building_xpu.yml new file mode 100644 index 00000000..cdebf297 --- /dev/null +++ b/.github/workflows/building_xpu.yml @@ -0,0 +1,86 @@ +name: Build XPU Wheels + +on: [workflow_call, workflow_dispatch] + +permissions: + contents: read + +jobs: + build_wheels: + runs-on: ubuntu-22.04 + environment: production + + strategy: + fail-fast: false + matrix: + python-version: ['3.10'] + torch-version: ['2.5.0', '2.6.0'] + xpu-version: ['xpu'] + + 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 + 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 Intel oneAPI DPC++ compiler + run: bash .github/workflows/xpu/Linux.sh + shell: bash + + - name: Install PyTorch ${{ matrix.torch-version }}+${{ matrix.xpu-version }} + run: | + pip install torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/${{ matrix.xpu-version }} + 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: 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"` + XPU_VERSION=`echo ${{ matrix.xpu-version }}` + echo "New version name: $VERSION+$TORCH_VERSION$XPU_VERSION" + sed -i "s/$VERSION/$VERSION+$TORCH_VERSION$XPU_VERSION/" gsplat/version.py + shell: bash + + - name: Upgrade pip + run: | + pip install --upgrade setuptools + pip install ninja pybind11 + shell: bash + + - name: Build wheel + run: | + pip install wheel + source /opt/intel/oneapi/setvars.sh + BUILD_SYCL=1 MAX_JOBS=$(nproc) python setup.py bdist_wheel --dist-dir=dist + shell: bash + + - 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 }}-ubuntu-22.04-${{ matrix.torch-version }}-${{ matrix.xpu-version }} + path: dist/*.whl diff --git a/.github/workflows/publish_xpu.yml b/.github/workflows/publish_xpu.yml new file mode 100644 index 00000000..a0d055f1 --- /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'] + 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..0e2e9ef9 --- /dev/null +++ b/.github/workflows/xpu/Linux.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +# Install Intel oneAPI DPC++ compiler for SYCL builds + +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-oneapi-dpcpp-cpp +sudo apt clean diff --git a/setup.py b/setup.py index 39269102..b18c0a4f 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ except (ImportError, AttributeError): pass -BUILD_SYCL = has_xpu +BUILD_SYCL = has_xpu or os.getenv("BUILD_SYCL", "0") == "1" BUILD_NO_CUDA = os.getenv("BUILD_NO_CUDA", "0") == "1" WITH_SYMBOLS = os.getenv("WITH_SYMBOLS", "0") == "1" LINE_INFO = os.getenv("LINE_INFO", "0") == "1" From 0305370440f7e03668ef9510f43d14222bb66f23 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 23 Feb 2026 19:45:52 +0000 Subject: [PATCH 45/56] Address PR review comments on SYCL/XPU wheel workflows Co-authored-by: ssheorey <41028320+ssheorey@users.noreply.github.com> --- .github/workflows/building_xpu.yml | 31 +++++++++++++++++------------- .github/workflows/publish_xpu.yml | 2 +- .github/workflows/xpu/Linux.sh | 8 ++++++-- 3 files changed, 25 insertions(+), 16 deletions(-) diff --git a/.github/workflows/building_xpu.yml b/.github/workflows/building_xpu.yml index cdebf297..823c26f9 100644 --- a/.github/workflows/building_xpu.yml +++ b/.github/workflows/building_xpu.yml @@ -13,9 +13,8 @@ jobs: strategy: fail-fast: false matrix: - python-version: ['3.10'] - torch-version: ['2.5.0', '2.6.0'] - xpu-version: ['xpu'] + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] + torch-version: ['2.10.0'] steps: - name: Checkout repository @@ -37,24 +36,30 @@ jobs: df -h shell: bash - - name: Install Intel oneAPI DPC++ compiler - run: bash .github/workflows/xpu/Linux.sh - shell: bash - - - name: Install PyTorch ${{ matrix.torch-version }}+${{ matrix.xpu-version }} + - name: Install PyTorch ${{ matrix.torch-version }}+xpu run: | - pip install torch==${{ matrix.torch-version }} --index-url https://download.pytorch.org/whl/${{ matrix.xpu-version }} + 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/Linux.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"` - XPU_VERSION=`echo ${{ matrix.xpu-version }}` - echo "New version name: $VERSION+$TORCH_VERSION$XPU_VERSION" - sed -i "s/$VERSION/$VERSION+$TORCH_VERSION$XPU_VERSION/" gsplat/version.py + echo "New version name: $VERSION+${TORCH_VERSION}xpu" + sed -i "s/$VERSION/$VERSION+${TORCH_VERSION}xpu/" gsplat/version.py shell: bash - name: Upgrade pip @@ -82,5 +87,5 @@ jobs: - uses: actions/upload-artifact@v4 with: # Include unique matrix values to avoid name collisions. - name: xpu_wheels_python${{ matrix.python-version }}-ubuntu-22.04-${{ matrix.torch-version }}-${{ matrix.xpu-version }} + name: xpu_wheels_python${{ matrix.python-version }}-ubuntu-22.04-${{ matrix.torch-version }} path: dist/*.whl diff --git a/.github/workflows/publish_xpu.yml b/.github/workflows/publish_xpu.yml index a0d055f1..986cf8a3 100644 --- a/.github/workflows/publish_xpu.yml +++ b/.github/workflows/publish_xpu.yml @@ -23,7 +23,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ['3.10'] + python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] steps: - name: Checkout code uses: actions/checkout@v4 diff --git a/.github/workflows/xpu/Linux.sh b/.github/workflows/xpu/Linux.sh index 0e2e9ef9..017a3344 100644 --- a/.github/workflows/xpu/Linux.sh +++ b/.github/workflows/xpu/Linux.sh @@ -1,6 +1,10 @@ #!/bin/bash -# Install Intel oneAPI DPC++ compiler for SYCL builds +# 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)'} wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB \ | gpg --dearmor \ @@ -10,5 +14,5 @@ echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt | sudo tee /etc/apt/sources.list.d/oneAPI.list sudo apt-get -qq update -sudo apt-get install -y intel-oneapi-dpcpp-cpp +sudo apt-get install -y intel-cpp-essentials-${VERSION} sudo apt clean From 2406dfbda1b7cd7162f5daf800c562d265f490bb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 23 Feb 2026 23:00:29 +0000 Subject: [PATCH 46/56] Initial plan From a0025c66a97fc7e9814ae24d509c1fb3c60abb12 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 23 Feb 2026 23:02:34 +0000 Subject: [PATCH 47/56] Fix XPU workflow: truncate oneAPI version X.Y.Z to X.Y for apt package install Co-authored-by: ssheorey <41028320+ssheorey@users.noreply.github.com> --- .github/workflows/xpu/Linux.sh | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/.github/workflows/xpu/Linux.sh b/.github/workflows/xpu/Linux.sh index 017a3344..eb6a3da2 100644 --- a/.github/workflows/xpu/Linux.sh +++ b/.github/workflows/xpu/Linux.sh @@ -6,6 +6,14 @@ 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 @@ -14,5 +22,5 @@ echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt | sudo tee /etc/apt/sources.list.d/oneAPI.list sudo apt-get -qq update -sudo apt-get install -y intel-cpp-essentials-${VERSION} +sudo apt-get install -y intel-cpp-essentials-${APT_VERSION} sudo apt clean From 1256a06f6de9a7b7247fe7df641586a64da13052 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Fri, 27 Feb 2026 11:41:21 -0800 Subject: [PATCH 48/56] Windows wheel workflow --- .github/workflows/building_xpu.yml | 41 ++++++++++--- .../workflows/generate_simple_index_pages.yml | 2 + .github/workflows/xpu/Windows.sh | 60 +++++++++++++++++++ 3 files changed, 95 insertions(+), 8 deletions(-) create mode 100644 .github/workflows/xpu/Windows.sh diff --git a/.github/workflows/building_xpu.yml b/.github/workflows/building_xpu.yml index 823c26f9..77f7e066 100644 --- a/.github/workflows/building_xpu.yml +++ b/.github/workflows/building_xpu.yml @@ -6,15 +6,35 @@ 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: ubuntu-22.04 + 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'] + torch-version: ['2.10'] steps: - name: Checkout repository @@ -28,6 +48,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Free up disk space + if: ${{ runner.os == 'Linux' }} run: | echo "Disk space before cleanup:" df -h @@ -51,7 +72,7 @@ jobs: exit 1 fi echo "Detected intel-sycl-rt version: ${ONEAPI_VERSION}" - bash .github/workflows/xpu/Linux.sh ${ONEAPI_VERSION} + bash .github/workflows/xpu/${RUNNER_OS}.sh ${ONEAPI_VERSION} shell: bash - name: Set version @@ -64,15 +85,19 @@ jobs: - name: Upgrade pip run: | - pip install --upgrade setuptools + pip install --upgrade setuptools wheel pip install ninja pybind11 shell: bash - name: Build wheel run: | - pip install wheel - source /opt/intel/oneapi/setvars.sh - BUILD_SYCL=1 MAX_JOBS=$(nproc) python setup.py bdist_wheel --dist-dir=dist + export MAX_JOBS=${NUMBER_OF_PROCESSORS:-$(nproc)} BUILD_SYCL=1 + if [ "${RUNNER_OS}" = "Windows" ]; then + cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && python setup.py bdist_wheel --dist-dir=dist' + else + source /opt/intel/oneapi/setvars.sh + python setup.py bdist_wheel --dist-dir=dist + fi shell: bash - name: Test wheel @@ -87,5 +112,5 @@ jobs: - uses: actions/upload-artifact@v4 with: # Include unique matrix values to avoid name collisions. - name: xpu_wheels_python${{ matrix.python-version }}-ubuntu-22.04-${{ matrix.torch-version }} + name: xpu_wheels_python${{ matrix.python-version }}-${{ matrix.os }}-${{ matrix.torch-version }} path: dist/*.whl 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/xpu/Windows.sh b/.github/workflows/xpu/Windows.sh new file mode 100644 index 00000000..2a051ef2 --- /dev/null +++ b/.github/workflows/xpu/Windows.sh @@ -0,0 +1,60 @@ +#!/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)'} + +if [[ ! "${VERSION}" =~ ^[0-9]+\.[0-9]+(\.[0-9]+)?$ ]]; then + echo "Error: VERSION '${VERSION}' does not match expected X.Y or X.Y.Z format." >&2 + exit 1 +fi + +# 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}" || "${INSTALLER_URL}" == "FILL_IN" ]]; 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 + oneDPL + oneTBB; override via ONEAPI_WINDOWS_COMPONENTS. +ONEAPI_WINDOWS_COMPONENTS="${ONEAPI_WINDOWS_COMPONENTS:-intel.oneapi.win.dpcpp-cpp-compiler;intel.oneapi.win.dpl;intel.oneapi.win.tbb.devel}" + +INSTALLER_FILE="/tmp/w_cpp-essentials_p_${VERSION}.exe" +echo "Downloading Intel C++ Essentials online installer from: ${INSTALLER_URL}" +curl -fL "${INSTALLER_URL}" --output "${INSTALLER_FILE}" + +echo "Installing components: ${ONEAPI_WINDOWS_COMPONENTS}" +PowerShell -NoProfile -Command "\$p = Start-Process -FilePath '${INSTALLER_FILE}' -ArgumentList '-s --action install --eula accept --components=${ONEAPI_WINDOWS_COMPONENTS}' -Wait -PassThru -NoNewWindow; exit \$p.ExitCode" +rm -f "${INSTALLER_FILE}" + +ONEAPI_SETVARS="/c/Program Files (x86)/Intel/oneAPI/setvars.bat" +if [[ ! -f "${ONEAPI_SETVARS}" ]]; then + echo "Error: oneAPI installation completed but setvars.bat was not found at '${ONEAPI_SETVARS}'." >&2 + exit 1 +fi + +echo "Verifying oneAPI environment..." +if ! cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && where icx'; then + echo "Error: icx (Intel DPC++/C++ Compiler) not found after sourcing setvars.bat." >&2 + exit 1 +fi + +echo "Intel oneAPI installation complete." From 1baa1728908a781d3dc675104b2a51ab4ac7e490 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Fri, 27 Feb 2026 12:30:00 -0800 Subject: [PATCH 49/56] Fix --- .github/workflows/xpu/Windows.sh | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/.github/workflows/xpu/Windows.sh b/.github/workflows/xpu/Windows.sh index 2a051ef2..83c25385 100644 --- a/.github/workflows/xpu/Windows.sh +++ b/.github/workflows/xpu/Windows.sh @@ -37,23 +37,16 @@ fi # Install only compiler + oneDPL + oneTBB; override via ONEAPI_WINDOWS_COMPONENTS. ONEAPI_WINDOWS_COMPONENTS="${ONEAPI_WINDOWS_COMPONENTS:-intel.oneapi.win.dpcpp-cpp-compiler;intel.oneapi.win.dpl;intel.oneapi.win.tbb.devel}" -INSTALLER_FILE="/tmp/w_cpp-essentials_p_${VERSION}.exe" +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}" echo "Installing components: ${ONEAPI_WINDOWS_COMPONENTS}" PowerShell -NoProfile -Command "\$p = Start-Process -FilePath '${INSTALLER_FILE}' -ArgumentList '-s --action install --eula accept --components=${ONEAPI_WINDOWS_COMPONENTS}' -Wait -PassThru -NoNewWindow; exit \$p.ExitCode" -rm -f "${INSTALLER_FILE}" - -ONEAPI_SETVARS="/c/Program Files (x86)/Intel/oneAPI/setvars.bat" -if [[ ! -f "${ONEAPI_SETVARS}" ]]; then - echo "Error: oneAPI installation completed but setvars.bat was not found at '${ONEAPI_SETVARS}'." >&2 - exit 1 -fi echo "Verifying oneAPI environment..." -if ! cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && where icx'; then - echo "Error: icx (Intel DPC++/C++ Compiler) not found after sourcing setvars.bat." >&2 +if ! cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && icx-cl --version'; then + echo "Error: icx-cl (Intel DPC++/C++ Compiler) not found after sourcing setvars.bat." >&2 exit 1 fi From 332882479c72cf790dc813003bbb8b1352aef180 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Fri, 27 Feb 2026 13:09:08 -0800 Subject: [PATCH 50/56] Fix setup.py for source only wheel --- setup.py | 104 ++++++++++++++++++++++++++----------------------------- 1 file changed, 49 insertions(+), 55 deletions(-) diff --git a/setup.py b/setup.py index b18c0a4f..6682dc85 100644 --- a/setup.py +++ b/setup.py @@ -4,33 +4,24 @@ import pathlib import platform import sys + from setuptools import find_packages, setup -import subprocess as sp __version__ = None exec(open("gsplat/version.py", "r").read()) URL = "https://github.com/nerfstudio-project/gsplat" -has_cuda = False +has_xpu = False try: import torch - has_cuda = torch.cuda.is_available() -except ImportError: + has_xpu = hasattr(torch, "xpu") and torch.xpu.is_available() +except (ImportError, AttributeError): pass -has_xpu = False -if not has_cuda: - 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" +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") @@ -41,48 +32,50 @@ print(f"Setting MAX_JOBS to {os.environ['MAX_JOBS']}") -from torch.utils.cpp_extension import BuildExtension - - -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 - - 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) - - cmake_args = [ - "cmake", - "-G", - "Ninja", - 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", "Release", "--", "-v", f"-j{jobs}"], - cwd=build_dir, - ) - - 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) + + cmake_args = [ + "cmake", + "-G", + "Ninja", + 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, + ) + cfg = "RelWithDebInfo" if WITH_SYMBOLS or LINE_INFO else "Release" + sp.check_call( + ["cmake", "--build", ".", "--config", cfg, "--", "-v", f"-j{jobs}"], + cwd=build_dir, + ) + return SyclBuildExtension def get_extensions(): @@ -168,12 +161,12 @@ def get_extensions(): if BUILD_SYCL: print("--- Configuring for SYCL build ---") - cmdclass = {"build_ext": SyclBuildExtension} + cmdclass = {"build_ext": get_ext()} ext_modules.append(Extension("gsplat.gsplat_sycl_kernels", sources=[])) elif not BUILD_NO_CUDA: print("--- Configuring for CUDA build ---") - ext_modules = get_extensions() cmdclass = {"build_ext": get_ext()} + ext_modules = get_extensions() else: print("--- Building without any C++/CUDA/SYCL extensions ---") @@ -194,6 +187,7 @@ def get_extensions(): "typing_extensions; python_version<'3.8'", ], extras_require={ + # dev dependencies. Install them by `pip install gsplat[dev]` "dev": [ "black[jupyter]==22.3.0", "isort==5.10.1", From b3a8f5dcc1bfed88efb4d66630ef213810837958 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Fri, 27 Feb 2026 16:21:13 -0800 Subject: [PATCH 51/56] installer fix --- .github/workflows/xpu/Windows.sh | 15 +++++---------- gsplat/sycl/CMakeLists.txt | 5 +++-- 2 files changed, 8 insertions(+), 12 deletions(-) diff --git a/.github/workflows/xpu/Windows.sh b/.github/workflows/xpu/Windows.sh index 83c25385..56d8bb5e 100644 --- a/.github/workflows/xpu/Windows.sh +++ b/.github/workflows/xpu/Windows.sh @@ -8,11 +8,6 @@ set -euo pipefail VERSION=${1:?'Usage: Windows.sh (e.g. Windows.sh 2025.3.1)'} -if [[ ! "${VERSION}" =~ ^[0-9]+\.[0-9]+(\.[0-9]+)?$ ]]; then - echo "Error: VERSION '${VERSION}' does not match expected X.Y or X.Y.Z format." >&2 - exit 1 -fi - # 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 @@ -27,22 +22,22 @@ declare -A INSTALLER_URLS=( ) INSTALLER_URL="${INSTALLER_URLS[${VERSION}]:-}" -if [[ -z "${INSTALLER_URL}" || "${INSTALLER_URL}" == "FILL_IN" ]]; then +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 + oneDPL + oneTBB; override via ONEAPI_WINDOWS_COMPONENTS. -ONEAPI_WINDOWS_COMPONENTS="${ONEAPI_WINDOWS_COMPONENTS:-intel.oneapi.win.dpcpp-cpp-compiler;intel.oneapi.win.dpl;intel.oneapi.win.tbb.devel}" - +# 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 '-s --action install --eula accept --components=${ONEAPI_WINDOWS_COMPONENTS}' -Wait -PassThru -NoNewWindow; exit \$p.ExitCode" +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" echo "Verifying oneAPI environment..." if ! cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && icx-cl --version'; then diff --git a/gsplat/sycl/CMakeLists.txt b/gsplat/sycl/CMakeLists.txt index b6a03e6b..4644f18e 100644 --- a/gsplat/sycl/CMakeLists.txt +++ b/gsplat/sycl/CMakeLists.txt @@ -27,8 +27,9 @@ except ImportError: sys.exit(1) lib_paths = cpp_extension.library_paths(True) -base_path = os.path.join(lib_paths[0], '${CMAKE_SHARED_LIBRARY_PREFIX}torch_python') -matches = glob.glob(base_path + '*') +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()}') From 02b390b044950e3ccd38805aa333c72f9ef42587 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Sat, 28 Feb 2026 09:21:55 -0800 Subject: [PATCH 52/56] fix --- .github/workflows/xpu/Windows.sh | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/.github/workflows/xpu/Windows.sh b/.github/workflows/xpu/Windows.sh index 56d8bb5e..d9e3e90d 100644 --- a/.github/workflows/xpu/Windows.sh +++ b/.github/workflows/xpu/Windows.sh @@ -39,10 +39,4 @@ curl -fL "${INSTALLER_URL}" --output "${INSTALLER_FILE}" 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" -echo "Verifying oneAPI environment..." -if ! cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && icx-cl --version'; then - echo "Error: icx-cl (Intel DPC++/C++ Compiler) not found after sourcing setvars.bat." >&2 - exit 1 -fi - -echo "Intel oneAPI installation complete." +rm "${INSTALLER_FILE}" \ No newline at end of file From 54de18ac3647c4430976f61ef6cb147ac4291648 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Sat, 28 Feb 2026 14:59:55 -0800 Subject: [PATCH 53/56] Separate Windows and Linux build steps. --- .github/workflows/building_xpu.yml | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/.github/workflows/building_xpu.yml b/.github/workflows/building_xpu.yml index 77f7e066..96290f12 100644 --- a/.github/workflows/building_xpu.yml +++ b/.github/workflows/building_xpu.yml @@ -89,16 +89,22 @@ jobs: pip install ninja pybind11 shell: bash - - name: Build wheel + - name: Build wheel (Windows) + if: ${{ runner.os == 'Windows' }} + shell: cmd run: | - export MAX_JOBS=${NUMBER_OF_PROCESSORS:-$(nproc)} BUILD_SYCL=1 - if [ "${RUNNER_OS}" = "Windows" ]; then - cmd //C '"C:\Program Files (x86)\Intel\oneAPI\setvars.bat" >nul && python setup.py bdist_wheel --dist-dir=dist' - else - source /opt/intel/oneapi/setvars.sh - python setup.py bdist_wheel --dist-dir=dist - fi + 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: | From 1442ff5a32d950bd6255cb918fae93f6585858d5 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey <41028320+ssheorey@users.noreply.github.com> Date: Sat, 28 Feb 2026 15:40:17 -0800 Subject: [PATCH 54/56] fix build type for windows --- .github/workflows/building_xpu.yml | 2 +- setup.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/building_xpu.yml b/.github/workflows/building_xpu.yml index 96290f12..f272c3c6 100644 --- a/.github/workflows/building_xpu.yml +++ b/.github/workflows/building_xpu.yml @@ -34,7 +34,7 @@ jobs: matrix: os: [ubuntu-22.04, windows-2022] python-version: ['3.10', '3.11', '3.12', '3.13', '3.14'] - torch-version: ['2.10'] + torch-version: ['2.10.0'] steps: - name: Checkout repository diff --git a/setup.py b/setup.py index 6682dc85..0a9a78e4 100644 --- a/setup.py +++ b/setup.py @@ -56,10 +56,12 @@ def run(self): 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, ] @@ -70,11 +72,11 @@ def run(self): cmake_args, cwd=build_dir, ) - cfg = "RelWithDebInfo" if WITH_SYMBOLS or LINE_INFO else "Release" sp.check_call( ["cmake", "--build", ".", "--config", cfg, "--", "-v", f"-j{jobs}"], cwd=build_dir, ) + return SyclBuildExtension From 21f1f8598495e6423e6e8cb6bbf78d62652fdf93 Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Thu, 12 Mar 2026 13:20:57 -0700 Subject: [PATCH 55/56] Update docs --- README.md | 2 ++ docs/Intel_XPU.md | 4 +++- examples/benchmarks/basic_4gpus.sh | 2 +- examples/requirements.txt | 2 +- examples/requirements_xpu.txt | 24 ++++++++++++++++++++++++ 5 files changed, 31 insertions(+), 3 deletions(-) create mode 100644 examples/requirements_xpu.txt 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 index 7fd7f09e..3de3bbb9 100644 --- a/docs/Intel_XPU.md +++ b/docs/Intel_XPU.md @@ -70,13 +70,15 @@ We evaluate gsplat-xpu on the Mip-NeRF 360 dataset and measure PSNR, SSIM, LPIPS ```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 - pip install --extra-index-url=https://download.pytorch.org/whl/xpu -r requirements.txt ``` 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 ``` 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/requirements.txt b/examples/requirements.txt index 73b3951e..a55535fb 100644 --- a/examples/requirements.txt +++ b/examples/requirements.txt @@ -20,5 +20,5 @@ tensorly pyyaml matplotlib git+https://github.com/rahul-goel/fused-ssim@b42e988db507702aa1198920ec7b36a5aa3f72b2 -#git+https://github.com/harry7557558/fused-bilagrid@90f9788e57d3545e3a033c1038bb9986549632fe +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 From 760d963d673ccad71f53b06444a5322447e63cce Mon Sep 17 00:00:00 2001 From: Sameer Sheorey Date: Thu, 12 Mar 2026 13:57:44 -0700 Subject: [PATCH 56/56] Add info about pre-built wheels. --- docs/Intel_XPU.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/Intel_XPU.md b/docs/Intel_XPU.md index 3de3bbb9..2a6b818f 100644 --- a/docs/Intel_XPU.md +++ b/docs/Intel_XPU.md @@ -25,6 +25,12 @@ The kernels are optimized and use mixed precision (some data is represented as h 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: