Normalmente, o código de simulação de física em tempo real é escrito em CUDA C++ de baixo nível para obter o máximo desempenho. Nesta publicação, apresentamos o NVIDIA Warp, um novo framework Python que facilita a criação de gráficos diferenciados e código de GPU de simulação em Python. O Warp fornece os componentes necessários para escrever código de simulação de alto desempenho, mas com a produtividade de trabalhar em uma linguagem interpretada como Python.
Ao final desta publicação, você aprenderá a usar o Warp para criar kernels CUDA em seu ambiente Python e aproveitar algumas das funcionalidades integradas de alto nível que facilitam a criação de simulações físicas complexas, como uma simulação de oceano (Figura 1).
Instalação
O Warp está disponível como uma biblioteca de código aberto do GitHub. Para fazer o download dos pacotes de lançamento e instalar em seu ambiente Python local, siga as instruções do README e use o seguinte comando:
pip install .
Inicialização
Após a importação, você deve inicializar explicitamente o Warp:
import warp as wp
wp.init()
Inicialização de kernels
O Warp usa o conceito de decoradores em Python para marcar funções que podem ser executadas na GPU. Por exemplo, você pode escrever um esquema simples de integração de partículas semi-implícitas da seguinte forma:
@wp.kernel
def integrate(x: wp.array(dtype=wp.vec3),
v: wp.array(dtype=wp.vec3),
f: wp.array(dtype=wp.vec3),
w: wp.array(dtype=float),
gravity: wp.vec3,
dt: float):
# thread id
tid = wp.tid()
x0 = x[tid]
v0 = v[tid]
# Semi-implicit Euler step
f_ext = f[tid] inv_mass = w[tid]
v1 = v0 + (f_ext * inv_mass + gravity) * dt
x1 = x0 + v1 * dt
# store results
x[tid] = x1
v[tid] = v1
Como o Warp é fortemente tipado, você deve fornecer dicas de tipo para os argumentos do kernel. Para abrir um kernel, use a seguinte sintaxe:
wp.launch(kernel=simple_kernel, # kernel to launch
dim=1024, # number of threads
inputs=[a, b, c], # parameters
device="cuda") # execution device
Ao contrário de frameworks baseados em tensores, como o NumPy, o Warp usa um modelo de programação baseado em kernel. A programação baseada em kernel corresponde mais de perto ao modelo de execução de GPU subjacente. Geralmente, é uma maneira mais natural de expressar o código de simulação que requer lógica condicional e operações de memória refinadas. No entanto, o Warp expõe esse modelo de programação centrado em threads de uma maneira fácil de usar que não requer conhecimento de baixo nível da arquitetura de GPU.
Modelo de compilação
O lançamento de um kernel ativa um pipeline de compilação Just-In-Time (JIT) que gera automaticamente o código do kernel C++/CUDA a partir de definições de funções do Python.
Todos os kernels pertencentes a um módulo Python são tempos de execução compilados em bibliotecas dinâmicas e PTX. A figura 2 mostra o pipeline de compilação, que envolve atravessar o AST de função e convertê-lo em código CUDA de linha reta que é, então, compilado e carregado de volta no processo Python.
O resultado desta compilação do JIT é armazenado em cache. Se a fonte do kernel de entrada não for alterada, os binários pré-combinados serão carregados em um modo de baixa sobrecarga.
Modelo de memória
As alocações de memória no Warp são expostas pelo tipo warp.array. Os arrays envolvem uma alocação de memória subjacente que pode estar na memória do host (CPU) ou do dispositivo (GPU). Ao contrário dos frameworks Tensor, os arrays no Warp são fortemente tipados e armazenam uma sequência linear de estruturas integradas (vec3, matrix33, quat, entre outros).
Você pode construir arrays a partir de listas Python ou arrays NumPy ou ainda inicializados usando uma sintaxe semelhante ao NumPy e ao PyTorch:
# allocate an uninitizalized array of vec3s
v = wp.empty(length=n, dtype=wp.vec3, device="cuda")
# allocate a zero-initialized array of quaternions
>q = wp.zeros(length=n, dtype=wp.quat, device="cuda")
# allocate and initialize an array from a numpy array
# will be automatically transferred to the specified device
v = wp.from_numpy(array, dtype=wp.vec3, device="cuda")
O Warp é compatível com os protocolos __array_interface__ e __cuda_array_interface__, que permitem visualizações de dados sem cópia entre frameworks baseados em Tensor. Por exemplo, para converter dados em NumPy, use o seguinte comando:
# automatically bring data from device back to host
view = device_array.numpy()
Recursos
O Warp inclui várias estruturas de dados de alto nível que facilitam a implementação de algoritmos de processamento de simulação e geometria.
Malhas
Malhas triangulares são onipresentes em simulação e computação gráfica. O Warp fornece um tipo integrado para gerenciar dados de malha que oferecem suporte para consultas geométricas, como pontos mais próximos, ray casting e verificações de sobreposição.
O exemplo a seguir mostra como usar o Warp para computar o ponto mais próximo em uma malha de um array de posições de entrada. Esse tipo de computação é o componente fundamental de muitos algoritmos na detecção de colisões (Figura 3). As consultas de malha do Warp simplificam a implementação desses métodos.
@wp.kernel
def project(positions: wp.array(dtype=wp.vec3),
mesh: wp.uint64,
output_pos: wp.array(dtype=wp.vec3),
output_face: wp.array(dtype=int)):
tid = wp.tid()
x = wp.load(positions, tid)
face_index = int(0)
face_u = float(0.0)
face_v = float(0.0)
sign = float(0.0)
max_dist = 2.0
if (wp.mesh_query_point(mesh, x, max_dist, sign, face_index, face_u, face_v)):
p = wp.mesh_eval_position(mesh, face_index, face_u, face_v)
output_pos[tid] = p
output_face[tid] = face_index
Volumes esparsos
Volumes esparsos são incrivelmente úteis para representar os dados da grade em domínios grandes, como campos de distância assinados (SDFs – Signed Distance Fields) para objetos ou velocidades complexas para fluxo de fluidos em grande escala. O Warp inclui suporte para volumes esparsos definidos usando o padrão NanoVDB. Construa volumes usando ferramentas OpenVDB padrão, como Blender, Houdini ou Maya, e depois faça a amostragem nos kernels do Warp.
Você pode criar volumes diretamente de arquivos de grade binária no disco ou na memória e, em seguida, experimentá-los usando a API de volumes:
wp.volume_sample_world(vol, xyz, mode) # world space sample using interpolation mode
wp.volume_sample_local(vol, uvw, mode) # volume space sample using interpolation mode
wp.volume_lookup(vol, ijk) # direct voxel lookup
wp.volume_transform(vol, xyz) # map point from voxel space to world space
wp.volume_transform_inv(vol, xyz) # map point from world space to volume space
Usando consultas de volume, você pode colidir com eficiência contra objetos complexos com sobrecarga de memória mínima.
Hash grids
Muitos métodos de simulação baseados em partículas, como o método dos elementos discretos (DEM – Discrete Element Method) ou hidrodinâmica das partículas suavizadas (SPH – Smoothed Particle Hydrodynamics), envolvem iterar sobre vizinhos espaciais para computar interações de força. As hash grids são uma estrutura de dados bem estabelecida para acelerar essas consultas de vizinhos mais próximos e são especialmente adequadas para a GPU.
As hash grids são construídas a partir de conjuntos de pontos da seguinte forma:
grid = wp.HashGrid(dim_x=128, dim_y=128, dim_z=128, device="cuda")
grid.build(points=p, radius=r)
Quando as hash grids são criadas, você pode consultá-las diretamente no código do kernel do usuário, conforme mostrado no exemplo a seguir, que computa a soma de todas as posições de partículas vizinhas:
@wp.kernel
def sum(grid : wp.uint64,
points: wp.array(dtype=wp.vec3),
output: wp.array(dtype=wp.vec3),
radius: float):
tid = wp.tid()
# query point
p = points[tid]
# create grid query around point
query = wp.hash_grid_query(grid, p, radius)
index = int(0)
sum = wp.vec3()
while(wp.hash_grid_query_next(query, index)):
neighbor = points[index]
# compute distance to neighbor point
dist = wp.length(p-neighbor)
if (dist <= radius):
sum += neighbor
output[tid] = sum
A figura 5 mostra um exemplo de simulação de materiais granulares de DEM para um material coeso. O uso da estrutura de dados integrada de hash grid permite escrever essa simulação em menos de 200 linhas de Python e a executar a taxas interativas para mais de 100 mil partículas.
O uso dos dados de hash grid do Warp permite avaliar facilmente as interações de força de paridade entre partículas vizinhas.
Diferencialidade
Os frameworks baseados em tensores, como o PyTorch e o JAX, oferecem gradientes de computação de tensores e são adequados para aplicações como treinamento de ML.
Um recurso exclusivo do Warp é a capacidade de gerar versões prospectivas e anteriores do código do kernel. Isso facilita a criação de simulações diferenciadas que podem propagar gradientes como parte de um pipeline de treinamento maior. Um cenário comum é usar frameworks de ML tradicionais para camadas de rede e o Warp para implementar camadas de simulação, permitindo uma diferenciabilidade de ponta a ponta.
Quando os gradientes forem necessários, você deve criar arrays com requires_grad=True. Por exemplo, a classe warp.Tape pode gravar as inicializações do kernel e reproduzi-las para calcular o gradiente de uma função de perda escalar em relação às entradas do kernel:
tape = wp.Tape()
# forward pass
with tape:
wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
wp.launch(kernel=compute2, inputs=[c, d], device="cuda")
wp.launch(kernel=loss, inputs=[d, l], device="cuda")
# reverse pass
tape.backward(loss=l)
Depois que o cálculo de atraso final for concluído, os gradientes em relação às entradas ficam disponíveis por um mapeamento no objeto Tape:
# gradient of loss with respect to input a
print(tape.gradients[a])
Resumo
Nesta publicação, apresentamos o NVIDIA Warp, um framework Python que facilita a criação de código de simulação diferenciado para a GPU. Encorajamos você a fazer download da versão prévia do Warp, compartilhar resultados e nos dar feedback.
Para obter mais informações, consulte os seguintes recursos: