Add batched CPU and CUDA mesh rasterization (#1697)

Summary:
Pull Request resolved: #1697

Add batched mesh rasterization for animation sequences that share topology, camera, and model transform. CUDA renders the full batch with one two-dimensional grid for each initialization, rasterization, and resolve stage, and uses smaller per-triangle blocks for the common small-triangle workload. CPU batches use `dispenso::parallel_for` after releasing the Python GIL.

Extend `pymomentum.torch.renderer.rasterize_mesh` to accept positions shaped `[batch, vertices, 3]`, return padded `[batch, height, width]` outputs, and preserve the existing unbatched API. Add a benchmark covering CPU, native CUDA batching, sequential CUDA calls, and stream-parallel calls.

For a 3,658-triangle character at 256x256 and batch size 128, native CUDA batching reduced total render time from 5.47 ms to 0.89 ms. Dispenso CPU batching reached 12.3k frames/s compared with 1.4k frames/s for a single-frame loop.

Reviewed By: yutingye

Differential Revision: D125535570

fbshipit-source-id: a563740a213f734ae6e39848379391b0c27bacfb

Author: cdtwigg

momentum/rasterizer/cuda_rasterizer.cpp

@@ -79,32 +79,42 @@ detail::CudaCamera makeCudaCamera(const Camera& camera, const Eigen::Matrix4f& m
 
 } // namespace
 
-void rasterizeMeshCuda(
+void rasterizeMeshesCuda(
     CudaRasterizerInputs inputs,
+    size_t batchSize,
     const Camera& camera,
     const Eigen::Matrix4f& modelMatrix,
     float nearClip,
-    CudaRasterizerBuffers buffers,
+    CudaBatchedRasterizerBuffers buffers,
     bool backfaceCulling,
     float depthOffset,
     cudaStream_t stream) {
   const auto* positionsModel = inputs.positionsModel;
   const size_t numVertices = inputs.numVertices;
   const auto* triangles = inputs.triangles;
   const size_t numTriangles = inputs.numTriangles;
+  MT_THROW_IF(batchSize == 0 || batchSize > 65535, "CUDA rasterizer batch size is out of range");
   MT_THROW_IF(nearClip <= 0.0f, "Near clip must be positive");
   MT_THROW_IF(
       camera.imageWidth() <= 0 || camera.imageHeight() <= 0, "Camera image must be non-empty");
   MT_THROW_IF(buffers.rowStride < camera.imageWidth(), "CUDA rasterizer row stride is too small");
-  MT_THROW_IF(buffers.visibilityBuffer == nullptr, "CUDA visibility buffer is null");
-  MT_THROW_IF(buffers.zBuffer == nullptr, "CUDA z buffer is null");
-  MT_THROW_IF(buffers.triangleIndexBuffer == nullptr, "CUDA triangle-index buffer is null");
   const auto imageHeight = static_cast<size_t>(camera.imageHeight());
   const auto rowStride = static_cast<size_t>(buffers.rowStride);
   MT_THROW_IF(
       rowStride > std::numeric_limits<size_t>::max() / imageHeight,
       "CUDA rasterizer output size overflows");
-  const size_t requiredBufferSize = imageHeight * rowStride;
+  const size_t imageRowsSize = imageHeight * rowStride;
+  MT_THROW_IF(
+      buffers.imageStride < 0 || static_cast<size_t>(buffers.imageStride) < imageRowsSize,
+      "CUDA rasterizer image stride is too small");
+  const auto imageStride = static_cast<size_t>(buffers.imageStride);
+  MT_THROW_IF(
+      batchSize - 1 > (std::numeric_limits<size_t>::max() - imageRowsSize) / imageStride,
+      "CUDA rasterizer output size overflows");
+  const size_t requiredBufferSize = (batchSize - 1) * imageStride + imageRowsSize;
+  MT_THROW_IF(buffers.visibilityBuffer == nullptr, "CUDA visibility buffer is null");
+  MT_THROW_IF(buffers.zBuffer == nullptr, "CUDA z buffer is null");
+  MT_THROW_IF(buffers.triangleIndexBuffer == nullptr, "CUDA triangle-index buffer is null");
   MT_THROW_IF(
       buffers.visibilityBufferSize < requiredBufferSize,
       "CUDA visibility buffer has {} elements; at least {} are required",
@@ -124,7 +134,12 @@ void rasterizeMeshCuda(
       numVertices > std::numeric_limits<size_t>::max() / 3 ||
           numTriangles > std::numeric_limits<size_t>::max() / 3,
       "CUDA mesh input size overflows");
-  const size_t requiredPositionElements = 3 * numVertices;
+  const size_t positionElementsPerMesh = 3 * numVertices;
+  MT_THROW_IF(
+      positionElementsPerMesh > 0 &&
+          batchSize > std::numeric_limits<size_t>::max() / positionElementsPerMesh,
+      "CUDA mesh input size overflows");
+  const size_t requiredPositionElements = batchSize * positionElementsPerMesh;
   const size_t requiredTriangleElements = 3 * numTriangles;
   MT_THROW_IF(numVertices > 0 && positionsModel == nullptr, "CUDA positions buffer is null");
   MT_THROW_IF(numTriangles > 0 && triangles == nullptr, "CUDA triangle buffer is null");
@@ -166,8 +181,9 @@ void rasterizeMeshCuda(
       device);
 
   throwIfCudaError(
-      detail::launchCudaRasterizeMesh(
+      detail::launchCudaRasterizeMeshes(
           positionsModel,
+          batchSize,
           numVertices,
           triangles,
           numTriangles,
@@ -176,11 +192,48 @@ void rasterizeMeshCuda(
           depthOffset,
           backfaceCulling,
           buffers.rowStride,
+          buffers.imageStride,
           buffers.visibilityBuffer,
           buffers.zBuffer,
           buffers.triangleIndexBuffer,
           stream),
-      "CUDA rasterizer launch");
+      "Batched CUDA rasterizer launch");
+}
+
+void rasterizeMeshCuda(
+    CudaRasterizerInputs inputs,
+    const Camera& camera,
+    const Eigen::Matrix4f& modelMatrix,
+    float nearClip,
+    CudaRasterizerBuffers buffers,
+    bool backfaceCulling,
+    float depthOffset,
+    cudaStream_t stream) {
+  const int64_t imageHeight = camera.imageHeight();
+  MT_THROW_IF(imageHeight <= 0 || buffers.rowStride < 0, "CUDA rasterizer output size is invalid");
+  MT_THROW_IF(
+      buffers.rowStride > std::numeric_limits<int64_t>::max() / imageHeight,
+      "CUDA rasterizer output size overflows");
+  const int64_t imageStride = imageHeight * buffers.rowStride;
+  rasterizeMeshesCuda(
+      inputs,
+      1,
+      camera,
+      modelMatrix,
+      nearClip,
+      CudaBatchedRasterizerBuffers{
+          .visibilityBuffer = buffers.visibilityBuffer,
+          .visibilityBufferSize = buffers.visibilityBufferSize,
+          .zBuffer = buffers.zBuffer,
+          .zBufferSize = buffers.zBufferSize,
+          .triangleIndexBuffer = buffers.triangleIndexBuffer,
+          .triangleIndexBufferSize = buffers.triangleIndexBufferSize,
+          .rowStride = buffers.rowStride,
+          .imageStride = imageStride,
+      },
+      backfaceCulling,
+      depthOffset,
+      stream);
 }
 
 } // namespace momentum::rasterizer

momentum/rasterizer/cuda_rasterizer.cu

@@ -12,6 +12,8 @@
 namespace momentum::rasterizer::detail {
 namespace {
 
+constexpr int32_t kRasterThreadsPerBlock = 32;
+
 struct TriangleSetup {
   float u[3];
   float v[3];
@@ -112,6 +114,7 @@ __global__ void initializeVisibilityKernel(
     int32_t imageWidth,
     int32_t imageHeight,
     int64_t rowStride,
+    int64_t imageStride,
     const float* zBuffer,
     uint64_t* visibilityBuffer) {
   const int64_t pixel = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
@@ -120,9 +123,10 @@ __global__ void initializeVisibilityKernel(
     return;
   }
 
+  const int64_t batchOffset = static_cast<int64_t>(blockIdx.y) * imageStride;
   const int32_t y = static_cast<int32_t>(pixel / imageWidth);
   const int32_t x = static_cast<int32_t>(pixel - static_cast<int64_t>(y) * imageWidth);
-  const int64_t offset = static_cast<int64_t>(y) * rowStride + x;
+  const int64_t offset = batchOffset + static_cast<int64_t>(y) * rowStride + x;
   const float existingDepth = zBuffer[offset];
   visibilityBuffer[offset] =
       isnan(existingDepth) ? 0 : static_cast<uint64_t>(floatToOrderedUint(existingDepth)) << 32;
@@ -138,8 +142,12 @@ __global__ void rasterizeVisibilityKernel(
     float depthOffset,
     bool backfaceCulling,
     int64_t rowStride,
+    int64_t imageStride,
     uint64_t* visibilityBuffer) {
+  const size_t batchIndex = blockIdx.y;
   const size_t triangleIndex = blockIdx.x;
+  positionsModel += 3 * batchIndex * numVertices;
+  visibilityBuffer += static_cast<int64_t>(batchIndex) * imageStride;
   if (triangleIndex >= numTriangles) {
     return;
   }
@@ -228,6 +236,7 @@ __global__ void resolveVisibilityKernel(
     int32_t imageWidth,
     int32_t imageHeight,
     int64_t rowStride,
+    int64_t imageStride,
     const uint64_t* visibilityBuffer,
     float* zBuffer,
     int32_t* triangleIndexBuffer) {
@@ -237,9 +246,10 @@ __global__ void resolveVisibilityKernel(
     return;
   }
 
+  const int64_t batchOffset = static_cast<int64_t>(blockIdx.y) * imageStride;
   const int32_t y = static_cast<int32_t>(pixel / imageWidth);
   const int32_t x = static_cast<int32_t>(pixel - static_cast<int64_t>(y) * imageWidth);
-  const int64_t offset = static_cast<int64_t>(y) * rowStride + x;
+  const int64_t offset = batchOffset + static_cast<int64_t>(y) * rowStride + x;
   const uint64_t key = visibilityBuffer[offset];
   const uint32_t triangleKey = static_cast<uint32_t>(key);
   if (triangleKey == 0) {
@@ -252,8 +262,9 @@ __global__ void resolveVisibilityKernel(
 
 } // namespace
 
-cudaError_t launchCudaRasterizeMesh(
+cudaError_t launchCudaRasterizeMeshes(
     const float* positionsModel,
+    size_t batchSize,
     size_t numVertices,
     const int32_t* triangles,
     size_t numTriangles,
@@ -262,24 +273,30 @@ cudaError_t launchCudaRasterizeMesh(
     float depthOffset,
     bool backfaceCulling,
     int64_t rowStride,
+    int64_t imageStride,
     uint64_t* visibilityBuffer,
     float* zBuffer,
     int32_t* triangleIndexBuffer,
     cudaStream_t stream) {
   const int64_t numPixels = static_cast<int64_t>(camera.imageWidth) * camera.imageHeight;
   const int32_t pixelBlocks = static_cast<int32_t>(
       (numPixels + kCudaRasterizerThreadsPerBlock - 1) / kCudaRasterizerThreadsPerBlock);
-
-  initializeVisibilityKernel<<<pixelBlocks, kCudaRasterizerThreadsPerBlock, 0, stream>>>(
-      camera.imageWidth, camera.imageHeight, rowStride, zBuffer, visibilityBuffer);
+  const uint32_t batchGridSize = static_cast<uint32_t>(batchSize);
+
+  initializeVisibilityKernel<<<
+      dim3(static_cast<uint32_t>(pixelBlocks), batchGridSize),
+      kCudaRasterizerThreadsPerBlock,
+      0,
+      stream>>>(
+      camera.imageWidth, camera.imageHeight, rowStride, imageStride, zBuffer, visibilityBuffer);
   if (const cudaError_t error = cudaGetLastError(); error != cudaSuccess) {
     return error;
   }
 
   if (numTriangles > 0) {
     rasterizeVisibilityKernel<<<
-        static_cast<uint32_t>(numTriangles),
-        kCudaRasterizerThreadsPerBlock,
+        dim3(static_cast<uint32_t>(numTriangles), batchGridSize),
+        kRasterThreadsPerBlock,
         0,
         stream>>>(
         positionsModel,
@@ -291,20 +308,58 @@ cudaError_t launchCudaRasterizeMesh(
         depthOffset,
         backfaceCulling,
         rowStride,
+        imageStride,
         visibilityBuffer);
     if (const cudaError_t error = cudaGetLastError(); error != cudaSuccess) {
       return error;
     }
   }
 
-  resolveVisibilityKernel<<<pixelBlocks, kCudaRasterizerThreadsPerBlock, 0, stream>>>(
+  resolveVisibilityKernel<<<
+      dim3(static_cast<uint32_t>(pixelBlocks), batchGridSize),
+      kCudaRasterizerThreadsPerBlock,
+      0,
+      stream>>>(
       camera.imageWidth,
       camera.imageHeight,
       rowStride,
+      imageStride,
       visibilityBuffer,
       zBuffer,
       triangleIndexBuffer);
   return cudaGetLastError();
 }
 
+cudaError_t launchCudaRasterizeMesh(
+    const float* positionsModel,
+    size_t numVertices,
+    const int32_t* triangles,
+    size_t numTriangles,
+    const CudaCamera& camera,
+    float nearClip,
+    float depthOffset,
+    bool backfaceCulling,
+    int64_t rowStride,
+    uint64_t* visibilityBuffer,
+    float* zBuffer,
+    int32_t* triangleIndexBuffer,
+    cudaStream_t stream) {
+  return launchCudaRasterizeMeshes(
+      positionsModel,
+      1,
+      numVertices,
+      triangles,
+      numTriangles,
+      camera,
+      nearClip,
+      depthOffset,
+      backfaceCulling,
+      rowStride,
+      static_cast<int64_t>(camera.imageHeight) * rowStride,
+      visibilityBuffer,
+      zBuffer,
+      triangleIndexBuffer,
+      stream);
+}
+
 } // namespace momentum::rasterizer::detail

momentum/rasterizer/cuda_rasterizer.h

@@ -45,6 +45,23 @@ struct CudaRasterizerBuffers {
   int64_t rowStride = 0;
 };
 
+/// Device buffers used by batched CUDA rasterization.
+///
+/// Each image uses rowStride elements per row and imageStride elements between batch entries. The
+/// corresponding size fields are element counts, not byte counts, and must cover the final image's
+/// rows from each buffer pointer. Existing zBuffer and triangleIndexBuffer values are preserved
+/// where no new triangle wins the depth test.
+struct CudaBatchedRasterizerBuffers {
+  uint64_t* visibilityBuffer = nullptr;
+  size_t visibilityBufferSize = 0;
+  float* zBuffer = nullptr;
+  size_t zBufferSize = 0;
+  int32_t* triangleIndexBuffer = nullptr;
+  size_t triangleIndexBufferSize = 0;
+  int64_t rowStride = 0;
+  int64_t imageStride = 0;
+};
+
 /// Rasterize a triangle mesh into CUDA-resident depth and triangle-index buffers.
 ///
 /// The input arrays must be contiguous CUDA buffers with shapes [numVertices, 3] and
@@ -70,4 +87,31 @@ void rasterizeMeshCuda(
     float depthOffset = 0.0f,
     cudaStream_t stream = nullptr);
 
+/// Rasterize a batch of triangle meshes with shared topology and camera.
+///
+/// inputs.positionsModel must be a contiguous CUDA buffer with shape
+/// [batchSize, inputs.numVertices, 3]. The triangle buffer, camera, and model matrix are shared by
+/// every batch entry. Work for the entire batch is launched asynchronously in three kernels on
+/// stream.
+///
+/// @param inputs CUDA input buffers, capacities, and mesh dimensions.
+/// @param batchSize Number of meshes; must be positive and at most 65535.
+/// @param camera Camera used for every batch entry.
+/// @param modelMatrix Additional model transform shared by every batch entry.
+/// @param nearClip Per-pixel near clipping distance; must be positive.
+/// @param buffers Batched CUDA output and scratch buffers.
+/// @param backfaceCulling Whether to discard back-facing triangles.
+/// @param depthOffset Value added to depth before the depth test and write.
+/// @param stream CUDA stream used for all launches.
+void rasterizeMeshesCuda(
+    CudaRasterizerInputs inputs,
+    size_t batchSize,
+    const Camera& camera,
+    const Eigen::Matrix4f& modelMatrix,
+    float nearClip,
+    CudaBatchedRasterizerBuffers buffers,
+    bool backfaceCulling = true,
+    float depthOffset = 0.0f,
+    cudaStream_t stream = nullptr);
+
 } // namespace momentum::rasterizer

momentum/rasterizer/cuda_rasterizer_kernels.cuh

@@ -16,12 +16,14 @@ namespace momentum::rasterizer::detail {
 
 inline constexpr int32_t kCudaRasterizerThreadsPerBlock = 256;
 
+/// Camera intrinsics models implemented by the CUDA projection path.
 enum class CudaCameraModel : int32_t {
   Pinhole,
   OpenCV,
   OpenCVFisheye,
 };
 
+/// Trivially copyable camera parameters passed to CUDA kernels.
 struct CudaCamera {
   CudaCameraModel model = CudaCameraModel::Pinhole;
   int32_t imageWidth = 0;
@@ -36,6 +38,7 @@ struct CudaCamera {
   float worldToEye[12] = {};
 };
 
+/// Launch the three CUDA kernels for one mesh.
 cudaError_t launchCudaRasterizeMesh(
     const float* positionsModel,
     size_t numVertices,
@@ -51,4 +54,22 @@ cudaError_t launchCudaRasterizeMesh(
     int32_t* triangleIndexBuffer,
     cudaStream_t stream);
 
+/// Launch the three CUDA kernels for a batch with shared topology and camera.
+cudaError_t launchCudaRasterizeMeshes(
+    const float* positionsModel,
+    size_t batchSize,
+    size_t numVertices,
+    const int32_t* triangles,
+    size_t numTriangles,
+    const CudaCamera& camera,
+    float nearClip,
+    float depthOffset,
+    bool backfaceCulling,
+    int64_t rowStride,
+    int64_t imageStride,
+    uint64_t* visibilityBuffer,
+    float* zBuffer,
+    int32_t* triangleIndexBuffer,
+    cudaStream_t stream);
+
 } // namespace momentum::rasterizer::detail