rasterize_points.cu

/*
 * Copyright (C) 2023, Inria
 * GRAPHDECO research group, https://team.inria.fr/graphdeco
 * All rights reserved.
 *
 * This software is free for non-commercial, research and evaluation use 
 * under the terms of the LICENSE.md file.
 *
 * For inquiries contact  george.drettakis@inria.fr
 */

#include <math.h>
#include <torch/extension.h>
#include <cstdio>
#include <sstream>
#include <iostream>
#include <tuple>
#include <stdio.h>
#include <cuda_runtime_api.h>
#include <memory>
#include <stdexcept>
#include "cuda_rasterizer/config.h"
#include "cuda_rasterizer/rasterizer.h"
#include <fstream>
#include <string>
#include <functional>

std::function<char*(size_t N)> resizeFunctional(torch::Tensor& t, const char* name = "buffer") {
    auto lambda = [&t, name](size_t N) {
        constexpr size_t kMaxReasonableBufferBytes = size_t(1) << 40;
        if (N > kMaxReasonableBufferBytes) {
            std::ostringstream oss;
            oss << "Refusing to allocate " << N << " bytes for CUDA rasterizer " << name << "; "
                << "this indicates a corrupted batch size or tile-overlap count.";
            throw std::runtime_error(oss.str());
        }
        t.resize_({(long long)N});
		return reinterpret_cast<char*>(t.contiguous().data_ptr());
    };
    return lambda;
}

std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
RasterizeGaussiansCUDA(
	const torch::Tensor& background,
	const torch::Tensor& means3D,
    const torch::Tensor& colors,
    const torch::Tensor& opacity,
	const torch::Tensor& scales,
	const torch::Tensor& rotations,
	const float scale_modifier,
	const torch::Tensor& cov3D_precomp,
	const torch::Tensor& viewmatrix,
	const torch::Tensor& projmatrix,
	const float tan_fovx, 
	const float tan_fovy,
    const int image_height,
    const int image_width,
	const torch::Tensor& sh,
	const int degree,
	const torch::Tensor& campos,
	const bool prefiltered,
	const bool antialiasing,
	const bool debug)
{
  if (means3D.ndimension() != 2 || means3D.size(1) != 3) {
    AT_ERROR("means3D must have dimensions (num_points, 3)");
  }
  
  const int P = means3D.size(0);
  const int H = image_height;
  const int W = image_width;

  auto int_opts = means3D.options().dtype(torch::kInt32);
  auto float_opts = means3D.options().dtype(torch::kFloat32);

  torch::Tensor out_color = torch::full({NUM_CHANNELS, H, W}, 0.0, float_opts);
  torch::Tensor out_invdepth = torch::full({0, H, W}, 0.0, float_opts);
  float* out_invdepthptr = nullptr;

  out_invdepth = torch::full({1, H, W}, 0.0, float_opts).contiguous();
  out_invdepthptr = out_invdepth.data<float>();

  torch::Tensor radii = torch::full({P}, 0, means3D.options().dtype(torch::kInt32));
  
  torch::Device device(torch::kCUDA);
  torch::TensorOptions options(torch::kByte);
  torch::Tensor geomBuffer = torch::empty({0}, options.device(device));
  torch::Tensor binningBuffer = torch::empty({0}, options.device(device));
  torch::Tensor imgBuffer = torch::empty({0}, options.device(device));
  std::function<char*(size_t)> geomFunc = resizeFunctional(geomBuffer, "geomBuffer");
  std::function<char*(size_t)> binningFunc = resizeFunctional(binningBuffer, "binningBuffer");
  std::function<char*(size_t)> imgFunc = resizeFunctional(imgBuffer, "imgBuffer");
  
  int rendered = 0;
  if(P != 0)
  {
	  int M = 0;
	  if(sh.size(0) != 0)
	  {
		M = sh.size(1);
      }

	  rendered = CudaRasterizer::Rasterizer::forward(
	    geomFunc,
		binningFunc,
		imgFunc,
	    P, degree, M,
		background.contiguous().data<float>(),
		W, H,
		means3D.contiguous().data<float>(),
		sh.contiguous().data_ptr<float>(),
		colors.contiguous().data<float>(), 
		opacity.contiguous().data<float>(), 
		scales.contiguous().data_ptr<float>(),
		scale_modifier,
		rotations.contiguous().data_ptr<float>(),
		cov3D_precomp.contiguous().data<float>(), 
		viewmatrix.contiguous().data<float>(), 
		projmatrix.contiguous().data<float>(),
		campos.contiguous().data<float>(),
		tan_fovx,
		tan_fovy,
		prefiltered,
		out_color.contiguous().data<float>(),
		out_invdepthptr,
		antialiasing,
		radii.contiguous().data<int>(),
		debug);
  }
  return std::make_tuple(rendered, out_color, radii, geomBuffer, binningBuffer, imgBuffer, out_invdepth);
}

std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 RasterizeGaussiansBackwardCUDA(
 	const torch::Tensor& background,
	const torch::Tensor& means3D,
	const torch::Tensor& radii,
    const torch::Tensor& colors,
	const torch::Tensor& opacities,
	const torch::Tensor& scales,
	const torch::Tensor& rotations,
	const float scale_modifier,
	const torch::Tensor& cov3D_precomp,
	const torch::Tensor& viewmatrix,
    const torch::Tensor& projmatrix,
	const float tan_fovx,
	const float tan_fovy,
    const torch::Tensor& dL_dout_color,
	const torch::Tensor& dL_dout_invdepth,
	const torch::Tensor& sh,
	const int degree,
	const torch::Tensor& campos,
	const torch::Tensor& geomBuffer,
	const int R,
	const torch::Tensor& binningBuffer,
	const torch::Tensor& imageBuffer,
	const bool antialiasing,
	const bool debug)
{
  const int P = means3D.size(0);
  const int H = dL_dout_color.size(1);
  const int W = dL_dout_color.size(2);
  
  int M = 0;
  if(sh.size(0) != 0)
  {	
	M = sh.size(1);
  }

  torch::Tensor dL_dmeans3D = torch::zeros({P, 3}, means3D.options());
  torch::Tensor dL_dmeans2D = torch::zeros({P, 3}, means3D.options());
  torch::Tensor dL_dcolors = torch::zeros({P, NUM_CHANNELS}, means3D.options());
  torch::Tensor dL_dconic = torch::zeros({P, 2, 2}, means3D.options());
  torch::Tensor dL_dopacity = torch::zeros({P, 1}, means3D.options());
  torch::Tensor dL_dcov3D = torch::zeros({P, 6}, means3D.options());
  torch::Tensor dL_dsh = torch::zeros({P, M, 3}, means3D.options());
  torch::Tensor dL_dscales = torch::zeros({P, 3}, means3D.options());
  torch::Tensor dL_drotations = torch::zeros({P, 4}, means3D.options());
  torch::Tensor dL_dinvdepths = torch::zeros({0, 1}, means3D.options());
  
  float* dL_dinvdepthsptr = nullptr;
  float* dL_dout_invdepthptr = nullptr;
  if(dL_dout_invdepth.size(0) != 0)
  {
	dL_dinvdepths = torch::zeros({P, 1}, means3D.options());
	dL_dinvdepths = dL_dinvdepths.contiguous();
	dL_dinvdepthsptr = dL_dinvdepths.data<float>();
	dL_dout_invdepthptr = dL_dout_invdepth.data<float>();
  }

  if(P != 0)
  {  
	  CudaRasterizer::Rasterizer::backward(P, degree, M, R,
	  background.contiguous().data<float>(),
	  W, H, 
	  means3D.contiguous().data<float>(),
	  sh.contiguous().data<float>(),
	  colors.contiguous().data<float>(),
	  opacities.contiguous().data<float>(),
	  scales.data_ptr<float>(),
	  scale_modifier,
	  rotations.data_ptr<float>(),
	  cov3D_precomp.contiguous().data<float>(),
	  viewmatrix.contiguous().data<float>(),
	  projmatrix.contiguous().data<float>(),
	  campos.contiguous().data<float>(),
	  tan_fovx,
	  tan_fovy,
	  radii.contiguous().data<int>(),
	  reinterpret_cast<char*>(geomBuffer.contiguous().data_ptr()),
	  reinterpret_cast<char*>(binningBuffer.contiguous().data_ptr()),
	  reinterpret_cast<char*>(imageBuffer.contiguous().data_ptr()),
	  dL_dout_color.contiguous().data<float>(),
	  dL_dout_invdepthptr,
	  dL_dmeans2D.contiguous().data<float>(),
	  dL_dconic.contiguous().data<float>(),  
	  dL_dopacity.contiguous().data<float>(),
	  dL_dcolors.contiguous().data<float>(),
	  dL_dinvdepthsptr,
	  dL_dmeans3D.contiguous().data<float>(),
	  dL_dcov3D.contiguous().data<float>(),
	  dL_dsh.contiguous().data<float>(),
	  dL_dscales.contiguous().data<float>(),
	  dL_drotations.contiguous().data<float>(),
	  antialiasing,
	  debug);
  }

  return std::make_tuple(dL_dmeans2D, dL_dcolors, dL_dopacity, dL_dmeans3D, dL_dcov3D, dL_dsh, dL_dscales, dL_drotations);
}

torch::Tensor markVisible(
		torch::Tensor& means3D,
		torch::Tensor& viewmatrix,
		torch::Tensor& projmatrix)
{ 
  const int P = means3D.size(0);
  
  torch::Tensor present = torch::full({P}, false, means3D.options().dtype(at::kBool));
 
  if(P != 0)
  {
	CudaRasterizer::Rasterizer::markVisible(P,
		means3D.contiguous().data<float>(),
		viewmatrix.contiguous().data<float>(),
		projmatrix.contiguous().data<float>(),
		present.contiguous().data<bool>());
  }
   
  return present;
}

// True batch kernel version: single forward call for N cameras using batch kernels
static std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
RasterizeGaussiansBatchKernelCUDAImpl(
	const torch::Tensor& background,
	const torch::Tensor& means3D,
    const torch::Tensor& colors,
    const torch::Tensor& opacity,
	const torch::Tensor& scales,
	const torch::Tensor& rotations,
	const float scale_modifier,
	const torch::Tensor& cov3D_precomp,
	const torch::Tensor& viewmatrix,
	const torch::Tensor& projmatrix,
	const float tan_fovx, 
	const float tan_fovy,
    const int image_height,
    const int image_width,
	const torch::Tensor& sh,
	const int degree,
	const torch::Tensor& campos,
	const bool prefiltered,
	const bool antialiasing,
	const bool debug,
	const bool render_color,
	const bool render_depth,
	const bool return_radii)
{
  if (!render_color && !render_depth) {
    AT_ERROR("Batch rasterization compact path must render color, depth, or both.");
  }
  if (means3D.ndimension() != 2 || means3D.size(1) != 3) {
    AT_ERROR("means3D must have dimensions (num_points, 3)");
  }
  
  // Check batch dimensions
  const int N = viewmatrix.size(0);  // number of cameras
  const int P = means3D.size(0);
  const int H = image_height;
  const int W = image_width;

  // Validate batch tensor shapes
  if (viewmatrix.size(1) != 4 || viewmatrix.size(2) != 4) {
    AT_ERROR("viewmatrix must have shape (N, 4, 4)");
  }
  if (projmatrix.size(1) != 4 || projmatrix.size(2) != 4) {
    AT_ERROR("projmatrix must have shape (N, 4, 4)");
  }
  if (campos.size(0) != N || campos.size(1) != 3) {
    AT_ERROR("campos must have shape (N, 3)");
  }
  if (N <= 0 || H <= 0 || W <= 0) {
    AT_ERROR("Batch rasterization expects positive N, image_height, and image_width.");
  }

  auto int_opts = means3D.options().dtype(torch::kInt32);
  auto float_opts = means3D.options().dtype(torch::kFloat32);

  // Output tensors: optional (N, 3, H, W), optional (N, 1, H, W), optional (N, P)
  torch::Tensor out_color = render_color
    ? torch::full({N, NUM_CHANNELS, H, W}, 0.0, float_opts)
    : torch::empty({0}, float_opts);
  torch::Tensor out_invdepth = render_depth
    ? torch::full({N, 1, H, W}, 0.0, float_opts)
    : torch::empty({0}, float_opts);
  torch::Tensor radii = return_radii
    ? torch::full({N, P}, 0, int_opts)
    : torch::empty({0}, int_opts);

  // Dummy buffers for backward compatibility (forward-only)
  torch::Device device(torch::kCUDA);
  torch::TensorOptions options(torch::kByte);
  torch::Tensor geomBuffer = torch::empty({0}, options.device(device));
  torch::Tensor binningBuffer = torch::empty({0}, options.device(device));
  torch::Tensor imgBuffer = torch::empty({0}, options.device(device));

  // Ensure input tensors are contiguous
  auto viewmatrix_c = viewmatrix.contiguous();
  auto projmatrix_c = projmatrix.contiguous();
  auto campos_c = campos.contiguous();
  auto means3D_c = means3D.contiguous();
  auto sh_c = sh.contiguous();
  auto colors_c = colors.contiguous();
  auto opacity_c = opacity.contiguous();
  auto scales_c = scales.contiguous();
  auto rotations_c = rotations.contiguous();
  auto cov3D_precomp_c = cov3D_precomp.contiguous();
  auto background_c = background.contiguous();

  const float* viewmatrix_ptr = viewmatrix_c.data_ptr<float>();
  const float* projmatrix_ptr = projmatrix_c.data_ptr<float>();
  const float* campos_ptr = campos_c.data_ptr<float>();
  
  const float* means3D_ptr = means3D_c.data_ptr<float>();
  const float* sh_ptr = sh_c.numel() > 0 ? sh_c.data_ptr<float>() : nullptr;
  const float* colors_ptr = colors_c.numel() > 0 ? colors_c.data_ptr<float>() : nullptr;
  const float* opacity_ptr = opacity_c.data_ptr<float>();
  const float* scales_ptr = scales_c.numel() > 0 ? scales_c.data_ptr<float>() : nullptr;
  const float* rotations_ptr = rotations_c.numel() > 0 ? rotations_c.data_ptr<float>() : nullptr;
  const float* cov3D_precomp_ptr = cov3D_precomp_c.numel() > 0 ? cov3D_precomp_c.data_ptr<float>() : nullptr;
  const float* bg_ptr = background_c.data_ptr<float>();

  if (render_color && ((sh_ptr == nullptr && colors_ptr == nullptr) || (sh_ptr != nullptr && colors_ptr != nullptr))) {
    AT_ERROR("Batch rasterization expects exactly one of SHs or precomputed colors when color output is enabled.");
  }

  float* color_ptr = render_color ? out_color.data_ptr<float>() : nullptr;
  float* depth_ptr = render_depth ? out_invdepth.data_ptr<float>() : nullptr;
  int* radii_ptr = return_radii ? radii.data_ptr<int>() : nullptr;

  int M = 0;
  if (sh.size(0) != 0) {
    M = sh.size(1);
  }

  // Call the batch forward kernel
  std::function<char*(size_t)> geomFunc = resizeFunctional(geomBuffer, "geomBuffer");
  std::function<char*(size_t)> binningFunc = resizeFunctional(binningBuffer, "binningBuffer");
  std::function<char*(size_t)> imgFunc = resizeFunctional(imgBuffer, "imgBuffer");

  int rendered = 0;
  if (P != 0) {
    rendered = CudaRasterizer::Rasterizer::forwardBatch(
      geomFunc,
      binningFunc,
      imgFunc,
      P, degree, M,
      N,
      bg_ptr,
      W, H,
      means3D_ptr,
      sh_ptr,
      colors_ptr,
      opacity_ptr,
      scales_ptr,
      scale_modifier,
      rotations_ptr,
      cov3D_precomp_ptr,
      viewmatrix_ptr,
      projmatrix_ptr,
      campos_ptr,
      tan_fovx,
      tan_fovy,
      prefiltered,
      color_ptr,
      depth_ptr,
      antialiasing,
      radii_ptr,
      debug);
  }

  if (debug) {
    auto err = cudaGetLastError();
    if (err != cudaSuccess) AT_ERROR("forwardBatch launch failed: ", cudaGetErrorString(err));
    err = cudaDeviceSynchronize();
    if (err != cudaSuccess) AT_ERROR("forwardBatch sync failed: ", cudaGetErrorString(err));
  }

  return std::make_tuple(rendered, out_color, radii, geomBuffer, binningBuffer, imgBuffer, out_invdepth);
}

std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
RasterizeGaussiansBatchKernelCUDA(
	const torch::Tensor& background,
	const torch::Tensor& means3D,
    const torch::Tensor& colors,
    const torch::Tensor& opacity,
	const torch::Tensor& scales,
	const torch::Tensor& rotations,
	const float scale_modifier,
	const torch::Tensor& cov3D_precomp,
	const torch::Tensor& viewmatrix,
	const torch::Tensor& projmatrix,
	const float tan_fovx, 
	const float tan_fovy,
    const int image_height,
    const int image_width,
	const torch::Tensor& sh,
	const int degree,
	const torch::Tensor& campos,
	const bool prefiltered,
	const bool antialiasing,
	const bool debug)
{
  return RasterizeGaussiansBatchKernelCUDAImpl(
    background, means3D, colors, opacity, scales, rotations, scale_modifier,
    cov3D_precomp, viewmatrix, projmatrix, tan_fovx, tan_fovy, image_height,
    image_width, sh, degree, campos, prefiltered, antialiasing, debug,
    true, true, true);
}

std::tuple<int, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
RasterizeGaussiansBatchKernelCompactCUDA(
	const torch::Tensor& background,
	const torch::Tensor& means3D,
    const torch::Tensor& colors,
    const torch::Tensor& opacity,
	const torch::Tensor& scales,
	const torch::Tensor& rotations,
	const float scale_modifier,
	const torch::Tensor& cov3D_precomp,
	const torch::Tensor& viewmatrix,
	const torch::Tensor& projmatrix,
	const float tan_fovx, 
	const float tan_fovy,
    const int image_height,
    const int image_width,
	const torch::Tensor& sh,
	const int degree,
	const torch::Tensor& campos,
	const bool prefiltered,
	const bool antialiasing,
	const bool debug,
	const bool render_color,
	const bool render_depth,
	const bool return_radii)
{
  return RasterizeGaussiansBatchKernelCUDAImpl(
    background, means3D, colors, opacity, scales, rotations, scale_modifier,
    cov3D_precomp, viewmatrix, projmatrix, tan_fovx, tan_fovy, image_height,
    image_width, sh, degree, campos, prefiltered, antialiasing, debug,
    render_color, render_depth, return_radii);
}