VkSplat: High-Performance 3DGS Training in Vulkan Compute

1. Declaração do Problema

Apesar da rápida adoção do Gaussian Splatting 3D (3DGS) para síntese de novas visualizações de alta qualidade, a implantação prática é dificultada por duas limitações principais:

• Gargalos de Desempenho: Pipelines de treinamento existentes são frequentemente subótimas em velocidade e uso de memória, particularmente para cenas grandes. Gargalos-chave incluem interseções de tiles falso-positivas durante a ordenação, contenção atômica durante a rasterização reversa e tratamento ineficiente de memória nos otimizadores (por exemplo, Adam do PyTorch).
• Bloqueio de Hardware: A maioria das implementações mais avançadas depende fortemente do ecossistema CUDA e do PyTorch, restringindo o treinamento a GPUs NVIDIA. Isso limita a acessibilidade para usuários com GPUs AMD, Intel ou móveis.
• Fragmentação: Tentativas anteriores de usar APIs de múltiplos fornecedores (como Vulkan) ou abordaram apenas a rasterização (não o treinamento de ponta a ponta) ou confiaram no PyTorch para o loop de treinamento, falhando em eliminar totalmente as dependências do CUDA.

2. Metodologia

Os autores apresentam o VkSplat, um pipeline de treinamento 3DGS totalmente de ponta a ponta implementado inteiramente em Vulkan Compute. O sistema é construído sobre a linguagem de sombreamento Slang (direcionada ao Vulkan) para garantir compatibilidade entre fornecedores, mantendo alto desempenho.

A metodologia central envolve uma reescrita completa do pipeline 3DGS com mudanças arquitetônicas específicas:

• Vulkan Compute Completo: Todo o pipeline, desde a inicialização e densificação até a otimização, é executado na GPU sem intervenção do lado da CPU ou dependências do PyTorch.
• Backend Slang: Aproveitar o Slang permite que o código vise múltiplos backends, embora a implementação atual se concentre no Vulkan.

3. Principais Contribuições Técnicas

O artigo detalha cinco grandes otimizações técnicas que permitem ao VkSplat superar as bases de referência baseadas em CUDA:

A. Eliminação Completa de Tiles via Interseção de Linha de Varredura

• Problema: Métodos existentes (por exemplo, [KKLD23], [RSP∗24]) geram interseções de tiles falso-positivas (verificando tiles que não realmente se sobrepõem ao Gaussiano), desperdiçando tempo de ordenação e rasterização.
• Solução: O VkSplat usa uma formulação de linha de varredura para calcular interseções exatas. Ao tratar o Gaussiano como uma elipse e calcular intervalos de forma fechada para linhas/colunas que se intersectam, o sistema gera uma lista precisa de pares tile-profundidade.
• Impacto: Elimina falso-positivos sem exigir uma passagem de eliminação separada, reduzindo o uso de VRAM e a sobrecarga de ordenação.

B. Rasterização Reversa Adaptativa

• Problema: Passadas reversas padrão por pixel causam alta contenção atômica ao atualizar gradientes do Gaussiano.
• Solução: Os autores implementam duas estratégias e usam um agendador de amostragem de Thompson para selecionar dinamicamente a mais rápida com base na cena:
  1. Paralelismo por Tile: Um bloco de threads por tile com tamanhos de lote de Gaussiano ajustados dinamicamente (otimizados para $S = \sqrt{NP}$) para minimizar a latência.
  2. Otimização de Memória Compartilhada: Uma abordagem de duas passagens (passagem direta computa transmitância na memória compartilhada; passagem reversa acumula gradientes por Gaussiano) para minimizar a divergência de threads.
• Impacto: Reduz significativamente a contenção atômica e melhora a utilização de threads.

C. Projeção Reversa e Otimizador Fundidos

• Problema: Implementações padrão armazenam gradientes separadamente e usam o Adam do PyTorch, levando a pegadas de memória desnecessárias e lançamentos de kernel.
• Solução: A passagem de projeção reversa e o otimizador Adam são fundidos em um único kernel.
  • Valores e gradientes para escalas e opacidades são transformados diretamente entre os espaços log/logit dentro do kernel.
  • Coeficientes de Harmônicos Esféricos (SH) são armazenados em um formato baseado em colunas alinhado com tamanhos de subgrupo (valores de 128 bits) para melhorar o coalescimento de memória.
• Impacto: Elimina o armazenamento intermediário de memória para gradientes e reduz a sobrecarga de lançamento de kernel.

D. Ordenação Tile-Profundidade de 32 bits

• Problema: A ordenação nativa do Vulkan é frequentemente mais lenta que os built-ins otimizados do CUDA.
• Solução: Em vez de usar chaves de 64 bits (ID de Tile de 32 bits + Profundidade de 32 bits), os autores usam chaves de 32 bits. Eles mapeiam a profundidade $z$ para um intervalo onde apenas os 23 bits de mantissa inferiores de um float FP32 variam, empacotando o ID do Tile nos bits superiores.
• Impacto: Permite ordenação eficiente de 32 bits em hardware Vulkan sem sacrificar precisão de profundidade ou qualidade de imagem.

E. Avaliação de Gradiente de Perda Totalmente Fundida

• Solução: Um único kernel computa a soma ponderada das perdas L1 e SSIM e seus gradientes diretamente.
• Otimização: Imagens de referência são armazenadas como UINT8 RGBA (evitando conversão FP32) e a máscara alfa é tratada com sobrecarga próxima de zero.
• Impacto: Remove etapas intermediárias de redução e conversões de layout de memória (canal-último para canal-primeiro).

4. Resultados

Os autores avaliaram o VkSplat em 7 cenas do conjunto de dados Mip-NeRF 360 contra a base de referência GSplat (uma implementação líder CUDA+PyTorch) em uma NVIDIA RTX 3090.

• Desempenho:
  • Velocidade: Alcançou um aceleramento de 3,3× sobre a base de referência GSplat (por exemplo, 412s vs. 1384s para densificação padrão).
  • VRAM: Reduziu o uso de memória em 33% (3,01 GiB vs. 4,56 GiB).
  • Granularidade: O VkSplat foi mais rápido em cada estágio individual do pipeline (Projeção, Tileamento, Rasterização, Perda, Otimizador).
• Qualidade:
  • Métricas de qualidade de imagem (PSNR, SSIM, LPIPS) foram idênticas à base de referência, confirmando que as otimizações não comprometeram a fidelidade.
• Compatibilidade entre Fornecedores:
  • Treinado com sucesso em uma AMD Radeon RX 7800 XT.
  • Embora a AMD tenha sido mais lenta no geral (1201s vs. 575s na NVIDIA), produziu métricas de qualidade e uso de VRAM consistentes. A lacuna de desempenho foi atribuída principalmente a transferências de memória host-dispositivo mais lentas na AMD, sugerindo espaço para ajustes específicos de hardware adicionais.

5. Significado

• Primeiro 3DGS Vulkan de Ponta a Ponta: Este é o primeiro pipeline de treinamento 3DGS totalmente baseado em Vulkan que alcança desempenho de última geração, quebrando o monopólio do CUDA.
• Agnosticismo de Hardware: Ao remover dependências do PyTorch e CUDA, o VkSplat habilita o treinamento 3DGS em GPUs AMD, Intel e móveis, democratizando o acesso à reconstrução 3D de alta qualidade.
• Liderança em Desempenho: Contra-intuitivamente, a implementação Vulkan supera a base de referência CUDA altamente otimizada, provando que otimizações cuidadosas de baixo nível de GPU (kernels fundidos, eliminação precisa, agendamento adaptativo) podem gerar ganhos mesmo em hardware não-NVIDIA.
• Extensibilidade: O uso do Slang sugere que o pipeline pode ser facilmente portado para outros backends como Metal, DirectX ou WebGPU, ampliando ainda mais o ecossistema.

Limitações: A implementação atual carece de algumas funcionalidades "prontas para produção" encontradas em treinadores maduros, como correção de exposição, supervisão de profundidade, loteamento e treinamento multi-GPU. No entanto, a arquitetura central foi projetada para acomodar essas funcionalidades.