Add Metal support (#31)

Author: pxl-th

Project.toml

@@ -1,9 +1,10 @@
 name = "GaussianSplatting"
 uuid = "991e6b22-92f0-46ac-8b70-8a93e7beee5d"
-version = "1.1.0"
+version = "1.2.0"
 authors = ["Anton Smirnov <[email protected]>"]
 
 [deps]
+AcceleratedKernels = "6a4ca0a5-0e36-4168-a932-d9be78d558f1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 CImGui = "5d785b6c-b76f-510e-a07c-3070796c7e87"
@@ -38,13 +39,16 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
 
 [extensions]
 GaussianSplattingAMDGPUExt = "AMDGPU"
 GaussianSplattingCUDAExt = "CUDA"
+GaussianSplattingMetalExt = "Metal"
 
 [compat]
 AMDGPU = "2"
+AcceleratedKernels = "0.4.3"
 Adapt = "4"
 BSON = "0.3"
 CImGui = "6"
@@ -61,10 +65,11 @@ ImageIO = "0.6"
 ImageMagick = "1.3"
 ImageTransformations = "0.10"
 KernelAbstractions = "0.9.34"
+Metal = "1.9"
 ModernGL = "1.1"
 NearestNeighbors = "0.4"
 NerfUtils = "0.2"
-NeuralGraphicsGL = "0.5"
+NeuralGraphicsGL = "0.5.1"
 PlyIO = "1.2"
 Quaternions = "0.7"
 Rotations = "1.7"

docs/src/index.md

@@ -7,8 +7,9 @@ Gaussian Splatting algorithm in pure Julia.
 ## Requirements
 
 - Julia 1.10 or higher.
-- AMD ([AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl)) or
-  Nvidia ([CUDA.jl](https://github.com/JuliaGPU/CUDA.jl)) capable machine.
+- [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl) or
+  [CUDA.jl](https://github.com/JuliaGPU/CUDA.jl) or
+  [Metal.jl](https://github.com/JuliaGPU/Metal.jl) capable machine.
 
 ## Install
 
@@ -24,14 +25,16 @@ GaussianSplatting.jl comes with a GUI application to train & view the gaussians.
 
 1. Add necessary packages:
    ```julia
-   ] add AMDGPU      # for AMD GPU
-   ] add CUDA # for Nvidia GPU
+   ] add AMDGPU # for AMD GPU
+   ] add CUDA   # for Nvidia GPU
+   ] add Metal  # for Apple GPU
    ```
 
 2. Run:
    ```julia
-   julia> using AMDGPU; kab = ROCBackend() # for AMD GPU
-   julia> using CUDA; kab = CUDABackend()  # for Nvidia GPU
+   julia> using AMDGPU; kab = ROCBackend()  # for AMD GPU
+   julia> using CUDA; kab = CUDABackend()   # for Nvidia GPU
+   julia> using Metal; kab = MetalBackend() # for Apple GPU
    julia> GaussianSplatting.gui(kab, "path-to-colmap-dataset-directory"; scale=1)
    ```
 

ext/GaussianSplattingAMDGPUExt/GaussianSplattingAMDGPUExt.jl

@@ -12,7 +12,7 @@ GaussianSplatting.base_array_type(::ROCBackend) = ROCArray
 
 GaussianSplatting.use_ak(::ROCBackend) = true
 
-function GaussianSplatting.allocate_pinned(kab, ::Type{T}, shape) where T
+function GaussianSplatting.allocate_pinned(::ROCBackend, ::Type{T}, shape) where T
     x = Array{T}(undef, shape)
     xd = unsafe_wrap(ROCArray, pointer(x), size(x))
     return x, xd

ext/GaussianSplattingMetalExt.jl

@@ -0,0 +1,19 @@
+module GaussianSplattingMetalExt
+
+using Metal
+using GaussianSplatting
+
+GaussianSplatting.base_array_type(::MetalBackend) = MtlArray
+
+GaussianSplatting.use_ak(::MetalBackend) = true
+
+function GaussianSplatting.allocate_pinned(::MetalBackend, ::Type{T}, shape) where T
+    xd = MtlArray{T, length(shape), Metal.SharedStorage}(undef, shape)
+    x = reshape(unsafe_wrap(Vector{T}, reshape(xd, :)), shape)
+    return x, xd
+end
+
+# Unregistered automatically in the array dtor.
+GaussianSplatting.unpin_memory(::MtlArray) = return
+
+end

src/GaussianSplatting.jl

@@ -31,6 +31,7 @@ using GLFW
 
 import CImGui.lib as iglib
 
+import AcceleratedKernels as AK
 import BSON
 import ChainRulesCore as CRC
 import ImageFiltering

src/gui/gui.jl

@@ -111,7 +111,7 @@ end
 
 # Viewer-only mode.
 function GSGUI(kab, gaussians::GaussianModel, camera::Camera; gl_kwargs...)
-    NGL.init(3, 0)
+    NGL.init(3, 2)
     context = NGL.Context("GaussianSplatting.jl"; gl_kwargs...)
     NGL.set_resize_callback!(context, resize_callback)
 
@@ -144,7 +144,7 @@ end
 
 # Training mode.
 function GSGUI(kab, dataset_path::String, scale::Int; gl_kwargs...)
-    NGL.init(3, 0)
+    NGL.init(3, 2)
     context = NGL.Context("GaussianSplatting.jl"; gl_kwargs...)
     NGL.set_resize_callback!(context, resize_callback)
 

src/rasterization/projection.jl

@@ -4,7 +4,7 @@ function project(
     rotations::AbstractMatrix{Float32};
     rast::GaussianRasterizer, camera::Camera,
     near_plane::Float32, far_plane::Float32,
-    radius_clip::Float32, blur_ϵ::Float32,
+    radius_clip::Int32, blur_ϵ::Float32,
 )
     (; width, height) = resolution(camera)
     @assert width % 16 == 0 && height % 16 == 0
@@ -57,7 +57,7 @@ function ∇project(
     conics::AbstractMatrix{Float32};
     rast::GaussianRasterizer, camera::Camera,
     near_plane::Float32, far_plane::Float32,
-    radius_clip::Float32, blur_ϵ::Float32,
+    radius_clip::Int32, blur_ϵ::Float32,
 )
     K = camera.intrinsics
     R_w2c = SMatrix{3, 3, Float32}(camera.w2c[1:3, 1:3])
@@ -110,7 +110,7 @@ function ChainRulesCore.rrule(::typeof(project),
 
     rast::GaussianRasterizer, camera::Camera,
     near_plane::Float32, far_plane::Float32,
-    radius_clip::Float32, blur_ϵ::Float32,
+    radius_clip::Int32, blur_ϵ::Float32,
 )
     means_2d, conics, compensations, depths = project(
         means_3d, scales, rotations;
@@ -151,7 +151,7 @@ end
     # Config.
     near_plane::Float32,
     far_plane::Float32,
-    radius_clip::Float32,
+    radius_clip::Int32,
     blur_ϵ::Float32,
 ) where {C <: Maybe{AbstractMatrix{Float32}}, RM}
     i = @index(Global)
@@ -194,10 +194,10 @@ end
 
     # Discard Gaussians outside of image plane.
     if (
-        (mean_2D[1] + radius) ≤ 0 ||
-        (mean_2D[1] - radius) ≥ resolution[1] ||
-        (mean_2D[2] + radius) ≤ 0 ||
-        (mean_2D[2] - radius) ≥ resolution[2]
+        (mean_2D[1] + radius) ≤ 0f0 ||
+        (mean_2D[1] - radius) ≥ Float32(resolution[1]) ||
+        (mean_2D[2] + radius) ≤ 0f0 ||
+        (mean_2D[2] - radius) ≥ Float32(resolution[2])
     )
         radii[i] = 0i32
         return

src/rasterization/rasterizer.jl

@@ -192,7 +192,7 @@ function (rast::GaussianRasterizer)(
         means_2d, conics, compensations, depths = project(
             means_3d, scales_act, rotations;
             rast, camera, near_plane=0.2f0, far_plane=1000f0,
-            radius_clip=3f0, blur_ϵ=0.3f0)
+            radius_clip=Int32(3), blur_ϵ=0.3f0)
 
         colors = spherical_harmonics(means_3d, shs; rast, camera, sh_degree)
 
@@ -252,7 +252,7 @@ function rasterize(
 
     # TODO make configurable.
     near_plane, far_plane = 0.2f0, 1000f0
-    radius_clip = 3f0 # In pixels.
+    radius_clip = Int32(3) # In pixels.
     blur_ϵ = 0.3f0
 
     project!(kab)(
@@ -315,7 +315,7 @@ function rasterize(
         rast.gstate.radii, rast.grid, BLOCK; ndrange=n)
 
     if use_ak(kab)
-        sortperm!(
+        AK.sortperm!(
             @view(rast.bstate.permutation[1:n_rendered]),
             @view(rast.bstate.gaussian_keys_unsorted[1:n_rendered]);
             temp=@view(rast.bstate.permutation_tmp[1:n_rendered]))

src/rasterization/utils.jl

@@ -13,8 +13,7 @@ sdiagm(x, y, z) = SMatrix{3, 3, Float32, 9}(
 
 gpu_floor(T, x) = unsafe_trunc(T, floor(x))
 gpu_ceil(T, x) = unsafe_trunc(T, ceil(x))
-
-gpu_cld(x, y::T) where T = (x + y - one(T)) ÷ y
+gpu_cld(x::X, y::T) where {X, T} = unsafe_trunc(T, floor(Float32(x + y - one(X)) / Float32(y)))
 
 Base.@propagate_inbounds function get_rect(
     pixel::SVector{2, Float32}, max_radius::Int32,
@@ -26,15 +25,6 @@ Base.@propagate_inbounds function get_rect(
     @inbounds rmax = SVector{2, Int32}(
         clamp(gpu_floor(Int32, gpu_cld(pixel[1] + max_radius, block[1])), 0i32, grid[1]),
         clamp(gpu_floor(Int32, gpu_cld(pixel[2] + max_radius, block[2])), 0i32, grid[2]))
-
-    # rblock = inv.(block)
-
-    # rmin = gpu_floor.(Int32, (pixel .- max_radius) .* rblock)
-    # rmin = clamp.(rmin, 0i32, grid)
-
-    # rmax = gpu_ceil.(Int32, (pixel .+ max_radius) .* rblock)
-    # # rmax = gpu_cld.(pixel .+ max_radius, block)
-    # rmax = clamp.(rmax, 0i32, grid)
     return rmin, rmax
 end