GaussianSplatDataPacking.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2024 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
 
// CPU data structures and packing helpers shared by the OpenGL and Metal
// Gaussian splat compute backends.  The structures mirror the std140/std430
// layouts expected by the GLSL/MSL compute shaders.
 
#pragma once
 
#include <Eigen/Core>
#include <Eigen/Geometry>
#include <cstdint>
#include <cstring>
#include <vector>
 
#include "open3d/visualization/rendering/Camera.h"
#include "open3d/visualization/rendering/gaussian_splat/GaussianSplatRenderer.h"
 
namespace open3d {
namespace visualization {
namespace rendering {
 
// ----- GPU-compatible structure definitions ----------------------------------
// These must be tightly packed to match the std140/std430 GLSL layouts.
 
struct Std430Vec4 {
    float x = 0.f, y = 0.f, z = 0.f, w = 0.f;
};
 
struct alignas(16) GaussianViewParams {
    // mat4 world_from_model (row-major, 4 vec4)
    float world_from_model[16];
    // mat4 view_from_world
    float view_from_world[16];
    // mat4 clip_from_view
    float clip_from_view[16];
    // vec4 camera_position_and_near
    float camera_position_and_near[4];
    // vec4 viewport_origin_and_size
    float viewport_origin_and_size[4];
    // uvec4 scene  (x=splat_count, y=gaussian_splat_sh_degree,
    // z=antialias_flag, w=screen_y_down_flag)
    std::uint32_t scene[4];
    // uvec4 tiles  (xy=tile_size, zw=tile_count)
    std::uint32_t tiles[4];
    // uvec4 limits (x=entry_capacity, y=max_tiles_per_splat,
    // z=max_tile_entries_total, w=tile_key_bits T = ceil(log2(tile_count)),
    // clamped to [1,31]: T bits for tile index, (32-T) bits for depth)
    std::uint32_t limits[4];
    // vec4 depth_range_and_flags (x=near, y=far, z=reserved, w=0)
    float depth_range_and_flags[4];
};
static_assert(sizeof(GaussianViewParams) == 288,
              "GaussianViewParams must be 288 bytes to match GLSL layout");
 
struct alignas(16) ProjectedComposite {
    std::uint32_t center_xy_fp16;
    float depth;
    float alpha;
    std::uint32_t packed_rgba8;
    alignas(16) Std430Vec4 inv_basis;
};
static_assert(sizeof(ProjectedComposite) == 32,
              "ProjectedComposite must be 32 bytes to match GLSL layout");
 
struct GaussianGpuCounters {
    std::uint32_t total_entries =
            0;                      
    std::uint32_t error_flags = 0;  
    std::uint32_t tile_count = 0;   
    std::uint32_t splat_count = 0;  
};
static_assert(sizeof(GaussianGpuCounters) == 4 * sizeof(std::uint32_t),
              "GaussianGpuCounters must be 16 bytes");
static constexpr std::size_t kGaussianCounterCount =
        sizeof(GaussianGpuCounters) / sizeof(std::uint32_t);  // 4
 
// Index constants for code that accesses the counters buffer as a raw uint32
// array rather than through the struct.
static constexpr std::size_t kGaussianCounterTotalEntriesIndex = 0;
static constexpr std::size_t kGaussianCounterErrorFlagsIndex = 1;
static constexpr std::size_t kGaussianCounterTileCountIndex = 2;
static constexpr std::size_t kGaussianCounterSplatCountIndex = 3;
 
inline constexpr std::uint32_t kGaussianGpuErrorTileEntryOverflow = 1u << 0;
inline constexpr std::uint32_t kGaussianGpuErrorSortCountClamped = 1u << 1;
inline constexpr std::uint32_t kGaussianGpuErrorMaxTilesPerSplatExceeded = 1u
                                                                           << 2;
inline constexpr std::uint32_t kGaussianGpuErrorKnownMask =
        kGaussianGpuErrorTileEntryOverflow | kGaussianGpuErrorSortCountClamped |
        kGaussianGpuErrorMaxTilesPerSplatExceeded;
 
// ----- CPU-side representation -----------------------------------------------
 
struct GaussianSplatPackedAttrs {
    std::vector<Std430Vec4> positions;  
    std::vector<std::uint32_t>
            scales;  
    std::vector<std::uint32_t>
            rotations;  
    std::vector<std::uint32_t> dc_opacity;  
    std::vector<std::uint32_t>
            sh_coefficients;  
    std::vector<std::uint32_t> visibility_mask;
    std::uint32_t splat_count = 0;
    int sh_degree = 0;  
    bool antialias = false;
};
 
struct GaussianSplatMergeItem {
    const GaussianSplatPackedAttrs* attrs = nullptr;
    bool visible = true;
};
 
struct PackedGaussianScene {
    GaussianViewParams view_params;
    std::uint32_t splat_count = 0;
    std::uint32_t tile_count = 0;
    std::uint32_t pixel_count = 0;
    bool valid = false;
};
 
// ----- GPU buffer sizing (absorbed from GaussianSplatBuffers) ----------------
 
inline constexpr std::uint32_t kGaussianRadixParamsStride = 256;
 
struct GaussianGpuBufferSizes {
    std::size_t projected_composite_size = 0;
    std::size_t tile_scalar_size = 0;
    std::size_t key_cap_size = 0;
    std::size_t histogram_buf_size = 0;
    std::size_t dispatch_args_size = 0;
    std::size_t radix_params_size = 0;
    std::uint32_t entries_capacity = 0;
    std::uint32_t radix_num_wg_cap = 0;
};
 
void ComputeGaussianGpuBufferSizes(const PackedGaussianScene& packed,
                                   GaussianGpuBufferSizes* out);
 
// ----- Packing helpers -------------------------------------------------------
 
PackedGaussianScene PackGaussianViewParams(
        const GaussianSplatPackedAttrs& attrs,
        const GaussianSplatRenderer::ViewRenderData& render_data,
        const GaussianSplatRenderer::RenderConfig& config);
 
void PackGaussianSplatAttrsDirect(const float* pts_ptr,
                                  std::size_t n,
                                  const float* scale_ptr,
                                  const float* rot_ptr,
                                  const float* f_dc_ptr,
                                  const float* opacity_ptr,
                                  const float* f_rest_ptr,
                                  int source_sh_degree,
                                  int desired_sh_degree,
                                  float min_opacity_logit,
                                  bool antialias,
                                  GaussianSplatPackedAttrs& out);
 
void MergeGaussianSplatPackedAttrs(
        const std::vector<GaussianSplatMergeItem>& items,
        GaussianSplatPackedAttrs* out,
        std::vector<std::uint32_t>* splat_starts = nullptr);
 
}  // namespace rendering
}  // namespace visualization
}  // namespace open3d