运行时参考#
本节介绍 Warp Python 运行时 API、如何管理内存、启动内核以及处理网格和体等对象的高级功能。本节描述的 API 旨在用于 Python 作用域并在 CPython 解释器中运行。有关 Kernel 作用域中可用函数的完整列表，请参阅内核参考部分。
内核#
内核是通过使用 @wp.kernel 装饰器进行标注的 Python 函数定义的。Python 函数的所有参数都必须使用其各自的类型进行标注。以下示例展示了一个将两个数组相加的简单内核
import warp as wp
@wp.kernel
def add_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float), c: wp.array(dtype=float)):
tid = wp.tid()
c[tid] = a[tid] + b[tid]
内核使用 wp.launch() 函数在特定设备（CPU/GPU）上启动。
wp.launch(add_kernel, dim=1024, inputs=[a, b, c], device="cuda")
请注意，所有内核输入都必须位于目标设备上，否则将引发运行时异常。内核可以使用多维网格边界启动。在这种情况下，线程不会被分配单个索引，而是在 n 维网格中的一个坐标，例如
wp.launch(complex_kernel, dim=(128, 128, 3), ...)
启动一个尺寸为 128 x 128 x 3 的 3D 线程网格。要获取每个线程的 3D 索引，请使用以下语法
i,j,k = wp.tid()
注意
目前，在 CPU 设备上启动的内核将串行执行。在 CUDA 设备上启动的内核将并行执行，并具有固定的块大小。
在 Warp 的编译模型中，内核使用 C++/CUDA 作为中间表示，即时编译成动态库和 PTX。为了避免内核代码在运行时过度重新编译，这些文件存储在一个以模块相关哈希命名的缓存目录中，以便重用之前编译的模块。Warp 初始化时会打印内核缓存的位置。wp.clear_kernel_cache() 可用于清除之前生成的编译产物的内核缓存，因为 Warp 不会自动尝试将缓存保持在某个大小以下。
- warp.kernel(f=None, *, enable_backward=None, module=None)[source]#
用于注册 Warp 内核的装饰器，该内核来自 Python 函数。函数必须为所有参数定义类型标注。函数不得返回任何内容。
示例
@wp.kernel def my_kernel(a: wp.array(dtype=float), b: wp.array(dtype=float)): tid = wp.tid() b[tid] = a[tid] + 1.0 @wp.kernel(enable_backward=False) def my_kernel_no_backward(a: wp.array(dtype=float, ndim=2), x: float): # the backward pass will not be generated i, j = wp.tid() a[i, j] = x @wp.kernel(module="unique") def my_kernel_unique_module(a: wp.array(dtype=float), b: wp.array(dtype=float)): # the kernel will be registered in new unique module created just for this # kernel and its dependent functions and structs tid = wp.tid() b[tid] = a[tid] + 1.0
- warp.launch(
- kernel,
- dim,
- inputs=[],
- outputs=[],
- adj_inputs=[],
- adj_outputs=[],
- device=None,
- stream=None,
- adjoint=False,
- record_tape=True,
- record_cmd=False,
- max_blocks=0,
- block_dim=256,
在目标设备上启动 Warp 内核
内核启动相对于调用 Python 线程是异步的。
- 参数：
kernel – Warp 内核函数的名称，使用
@wp.kernel装饰器装饰inputs (Sequence) – 内核的输入参数（可选）
outputs (Sequence) – 输出参数（可选）
adj_inputs (Sequence) – 伴随输入（可选）
adj_outputs (Sequence) – 伴随输出（可选）
stream (Stream | None) – 要启动的流。
adjoint (bool) – 是运行前向传播还是反向传播（通常使用
False）。record_cmd (bool) – 当为
True时，启动将返回一个Launch对象。实际启动直到用户调用Launch.launch()方法时才会发生。max_blocks (int) – 要使用的 CUDA 线程块的最大数量。仅对 CUDA 内核启动有效。如果为负或零，则将使用最大硬件值。
block_dim (int) – 每个块的线程数（对于“cpu”设备始终为 1）。
- warp.launch_tiled(*args, **kwargs)[source]#
一个辅助方法，用于启动一个网格，该网格具有一个额外的尾部维度，其大小等于块大小。
例如，要启动一个 2D 网格，其中每个元素分配有 64 个线程，您将使用以下方法
wp.launch_tiled(kernel, [M, N], inputs=[...], block_dim=64)
这等效于以下内容
wp.launch(kernel, [M, N, 64], inputs=[...], block_dim=64)
在您的内核代码中，您可以像往常一样获取线程的前两个索引，如果需要，可以忽略隐式的第三个维度
@wp.kernel def compute() i, j = wp.tid() ...
- warp.clear_kernel_cache()[source]#
清除内核缓存目录中先前生成的源代码和编译器工件。
只有以
wp_开头的目录将被删除。此函数仅清除当前 Warp 版本的缓存。LTO 工件不受影响。- 返回类型：
None
运行时内核创建#
Warp 允许即时生成内核，并提供各种定制，包括闭包支持。有关最新功能，请参阅代码生成部分。
启动对象#
Launch 对象是减少多次启动内核开销的一种方式。Launch 对象是通过调用 wp.launch() 并设置 record_cmd=True 返回的。这会存储启动内核所需的各种开销操作的结果，但会延迟实际的内核启动，直到调用 Launch.launch() 方法。
与图不同，Launch 对象不会减少为在 GPU 上执行内核而准备驱动程序的开销。另一方面，Launch 对象没有 CUDA 图的存储和初始化开销，并且允许使用 Launch.set_dim() 修改启动维度，以及使用 Launch.set_params() 和 Launch.set_param_by_name() 等函数修改内核参数。此外，Launch 对象也可用于减少在 CPU 上运行内核的开销。
- class warp.Launch(
- kernel,
- device,
- hooks=None,
- params=None,
- params_addr=None,
- bounds=None,
- max_blocks=0,
- block_dim=256,
- adjoint=False,
表示内核启动所需的所有数据，以便快速重放启动。
用户不应直接实例化此类，而是使用
wp.launch(..., record_cmd=True)来记录启动。- 参数：
device (Device)
hooks (Optional[KernelHooks])
params (Optional[Sequence[Any]])
params_addr (Optional[Sequence[ctypes.c_void_p]])
bounds (Optional[launch_bounds_t])
max_blocks (int)
block_dim (int)
adjoint (bool)
数组#
数组是 Warp 中最基本的内存抽象。可以通过以下全局构造函数创建数组
wp.empty(shape=1024, dtype=wp.vec3, device="cpu")
wp.zeros(shape=1024, dtype=float, device="cuda")
wp.full(shape=1024, value=10, dtype=int, device="cuda")
数组也可以直接从 NumPy ndarrays 构建，如下所示
r = np.random.rand(1024)
# copy to Warp owned array
a = wp.array(r, dtype=float, device="cpu")
# return a Warp array wrapper around the NumPy data (zero-copy)
a = wp.array(r, dtype=float, copy=False, device="cpu")
# return a Warp copy of the array data on the GPU
a = wp.array(r, dtype=float, device="cuda")
请注意，对于多维数据，必须明确指定 dtype 参数，例如
r = np.random.rand((1024, 3))
# initialize as an array of vec3 objects
a = wp.array(r, dtype=wp.vec3, device="cuda")
如果形状不兼容，将引发错误。
Warp 数组也可以从定义了 __cuda_array_interface__ 属性的对象构建。例如
import cupy
import warp as wp
device = wp.get_cuda_device()
r = cupy.arange(10)
# return a Warp array wrapper around the cupy data (zero-copy)
a = wp.array(r, device=device)
可以使用 array.to() 在设备之间移动数组
host_array = wp.array(a, dtype=float, device="cpu")
# allocate and copy to GPU
device_array = host_array.to("cuda")
此外，可以使用 wp.copy() 在不同内存空间的数组之间复制数据
src_array = wp.array(a, dtype=float, device="cpu")
dest_array = wp.empty_like(host_array)
# copy from source CPU buffer to GPU
wp.copy(dest_array, src_array)
- class warp.array(*args, **kwargs)[source]#
一个包含相同类型值的固定大小多维数组。
- dtype#
数组的数据类型。
- 类型：
DType
- __init__(
- data=None,
- dtype=Any,
- shape=None,
- strides=None,
- length=None,
- ptr=None,
- capacity=None,
- device=None,
- pinned=False,
- copy=True,
- owner=False,
- deleter=None,
- ndim=None,
- grad=None,
- requires_grad=False,
构造一个新的 Warp 数组对象
当
data参数是有效的 list、tuple 或 ndarray 时，数组将从该对象的数据构建。对于未在内存中连续存储的对象（例如：list），数据将首先展平，然后再传输到 device 指定的内存空间。第二种构造路径发生在
ptr参数是一个非零的 uint64 值，表示现有数组数据所在的内存起始地址，例如：来自外部库或 C 库。内存分配应位于 device 参数指定的同一设备上，用户应适当设置 length 和 dtype 参数。如果既未指定
data也未指定ptr，则接下来检查shape或length参数。这种构造路径可用于创建新的未初始化的数组，但建议用户改用wp.empty()、wp.zeros()或wp.full()来创建新数组。如果以上所有参数均未指定，则会构造一个简单的类型注解。这用于注解内核参数或结构体成员（例如，``arr: wp.array(dtype=float)``）。在这种情况下，只考虑
dtype和ndim，并且不为数组分配内存。- 参数：
data (List | Tuple | ndarray[tuple[int, ...], dtype[_ScalarType_co]] | None) – 用于构造数组的对象，可以是 Tuple、List，或通常是任何可转换为 np.array 的类型
dtype (Any) – 其中一种可用的数据类型，例如
warp.float32、warp.mat33，或自定义的结构体。如果 dtype 为Any且 data 是 ndarray，则会从数组数据类型推断。length (int | None) – 数据类型的元素数量（已弃用，用户应使用
shape参数）。ptr (int | None) – 要别名的外部内存地址的地址（
data应为None）。capacity (int | None) –
ptr分配的最大字节大小（data应为None）。device (Devicelike) – 数组所在的设备。
copy (bool) – 传入的
data是会被复制还是别名。别名要求传入的data已位于指定的device上且数据类型匹配。owner (bool) – 数组在删除时是否会尝试释放底层内存（已弃用，如果您希望将所有权转移给 Warp，请传递
deleter）。deleter (Callable[[int, int], None] | None) – 数组删除时调用的函数，接受两个参数：指针和大小。
grad (array | None) – 在反向传播中累积梯度的数组。如果为
None且requires_grad为True，则会自动分配一个梯度数组。pinned (bool) – 是否分配固定（pinned）主机内存，这允许异步主机-设备传输（仅适用于
device="cpu"）。ndim (int | None)
- fill_(value)[source]#
将所有数组条目设置为 value。
- 参数：
value – 要设置的数组中每个条目的值。必须可转换为数组的
dtype。- Raises:
ValueError – 如果 value 无法转换为数组的
dtype。
Examples
填充向量或矩阵数组时，
fill_()可以接受列表或其他序列。>>> arr = wp.zeros(2, dtype=wp.mat22) >>> arr.numpy() array([[[0., 0.], [0., 0.]], [[0., 0.], [0., 0.]]], dtype=float32) >>> arr.fill_([[1, 2], [3, 4]]) >>> arr.numpy() array([[[1., 2.], [3., 4.]], [[1., 2.], [3., 4.]]], dtype=float32)
- assign(src)[source]#
如果
src尚未是warp.array，则将其包装在其中，并将内容复制到self。
- numpy()[source]#
将数组转换为
numpy.ndarray（通过数组接口协议别名内存）。如果数组在 GPU 上，将自动执行同步的设备到主机复制（在 CUDA 默认流上），以确保所有未完成的工作都已完成。
- cptr()[source]#
返回数组地址的 ctypes 转换。
Notes
只有 CPU 数组支持此方法。
数组必须是连续的。
对该对象的访问不进行边界检查。
对于
float16类型，返回指向内部uint16表示的指针。
- transpose(axes=None)[source]#
返回一个经过轴转置的数组的零复制视图。
注意：转置操作将返回一个访问模式非连续的数组。
- 参数：
axes (可选) – 指定轴的排列方式。如果未指定，轴顺序将反转。
- ipc_handle()[source]#
将数组的 IPC handle 返回为 64 字节的
bytes对象。在另一个进程中，可以使用此 handle 通过
from_ipc_handle()获取一个共享同一底层内存分配的array。IPC 目前仅在 Linux 上支持。此外，IPC 仅支持使用默认内存分配器分配的数组。
使用参数
interprocess=True创建的Event对象也可以在进程间共享，以同步 GPU 工作。示例
临时使用默认内存分配器分配一个数组并获取其 IPC handle。
with wp.ScopedMempool("cuda:0", False): test_array = wp.full(1024, value=42.0, dtype=wp.float32, device="cuda:0") ipc_handle = test_array.ipc_handle()
- Raises:
RuntimeError – 数组未关联 CUDA 设备。
RuntimeError – CUDA 设备似乎不支持 IPC。
RuntimeError – 数组是使用 mempool 内存分配器分配的。
- 返回类型：
多维数组#
通过传递一个元组来指定每个维度的大小，可以构造最多四维的多维数组。
以下构造了一个大小为 1024 x 16 的二维数组。
wp.zeros(shape=(1024, 16), dtype=float, device="cuda")
将多维数组传递给内核时，用户必须在内核签名中指定预期的数组维度，例如，要将二维数组传递给内核，使用 ndim=2 参数指定维数。
@wp.kernel
def test(input: wp.array(dtype=float, ndim=2)):
提供了常用数组大小的类型提示助手，例如：array2d()、array3d()，它们等同于调用 array(..., ndim=2)` 等。要索引多维数组，请使用以下内核语法。
# returns a float from the 2d array
value = input[i,j]
要创建数组切片，请使用以下语法，其中索引数量小于数组维度。
# returns an 1d array slice representing a row of the 2d array
row = input[i]
切片操作符可以串联，例如：s = array[i][j][k]。切片可以传递给 wp.func 用户函数，前提是该函数也声明了预期的数组维度。目前，仅支持单索引切片。
提供了以下构造方法用于分配零初始化和空（未初始化）数组。
- warp.zeros(
- shape=None,
- dtype=float,
- device=None,
- requires_grad=False,
- pinned=False,
- **kwargs,
返回一个零初始化的数组。
- warp.zeros_like(src, device=None, requires_grad=None, pinned=None)[source]#
返回一个与另一个数组类型和维度相同的零初始化数组。
- warp.ones(
- shape=None,
- dtype=float,
- device=None,
- requires_grad=False,
- pinned=False,
- **kwargs,
返回一个全一初始化的数组。
- warp.ones_like(src, device=None, requires_grad=None, pinned=None)[source]#
返回一个与另一个数组类型和维度相同的全一初始化数组。
- warp.full(
- shape=None,
- value=0,
- dtype=Any,
- device=None,
- requires_grad=False,
- pinned=False,
- **kwargs,
返回一个所有元素都初始化为给定值的数组。
- warp.full_like(src, value, device=None, requires_grad=None, pinned=None)[source]#
返回一个与另一个数组类型和维度相同、所有元素都初始化为给定值的数组。
- warp.empty(
- shape=None,
- dtype=float,
- device=None,
- requires_grad=False,
- pinned=False,
- **kwargs,
返回一个未初始化的数组。
- warp.empty_like(src, device=None, requires_grad=None, pinned=None)[source]#
返回一个与另一个数组类型和维度相同的未初始化数组。
- warp.copy(dest, src, dest_offset=0, src_offset=0, count=0, stream=None)[source]#
将数组内容从 src 复制到 dest。
- 参数：
如果指定了流，则可以来自任何设备。如果省略了流，Warp 将根据以下规则选择一个流：(1) 如果目标数组位于 CUDA 设备上，则使用目标设备上的当前流。(2) 否则，如果源数组位于 CUDA 设备上，则使用源设备上的当前流。
如果源或目标都不在 CUDA 设备上，则复制时不使用流。
数据类型#
标量类型#
数组结构支持以下标量存储类型
bool |
布尔型 |
int8 |
有符号字节 |
uint8 |
无符号字节 |
int16 |
有符号短整型 |
uint16 |
无符号短整型 |
int32 |
有符号整型 |
uint32 |
无符号整型 |
int64 |
有符号长整型 |
uint64 |
无符号长整型 |
float16 |
半精度浮点数 |
float32 |
单精度浮点数 |
float64 |
双精度浮点数 |
Warp 支持使用 float 和 int 作为 wp.float32 和 wp.int32 的别名。
向量#
Warp 提供了用于常见模拟和图形问题的内置数学和几何类型。这些类型的运算符和函数的完整参考可以在 内核参考 中找到。
Warp 支持任意长度/数值类型的向量。内置的具体类型如下
vec2 vec3 vec4 |
2D、3D、4D 单精度浮点向量 |
vec2b vec3b vec4b |
2D、3D、4D 有符号字节向量 |
vec2ub vec3ub vec4ub |
2D、3D、4D 无符号字节向量 |
vec2s vec3s vec4s |
2D、3D、4D 有符号短整型向量 |
vec2us vec3us vec4us |
2D、3D、4D 无符号短整型向量 |
vec2i vec3i vec4i |
2D、3D、4D 有符号整型向量 |
vec2ui vec3ui vec4ui |
2D、3D、4D 无符号整型向量 |
vec2l vec3l vec4l |
2D、3D、4D 有符号长整型向量 |
vec2ul vec3ul vec4ul |
2D、3D、4D 无符号长整型向量 |
vec2h vec3h vec4h |
2D、3D、4D 半精度浮点向量 |
vec2f vec3f vec4f |
2D、3D、4D 单精度浮点向量 |
vec2d vec3d vec4d |
2D、3D、4D 双精度浮点向量 |
spatial_vector |
6D 单精度浮点向量 |
spatial_vectorf |
6D 单精度浮点向量 |
spatial_vectord |
6D 双精度浮点向量 |
spatial_vectorh |
6D 半精度浮点向量 |
向量支持大多数标准线性代数运算，例如
@wp.kernel
def compute( ... ):
# basis vectors
a = wp.vec3(1.0, 0.0, 0.0)
b = wp.vec3(0.0, 1.0, 0.0)
# take the cross product
c = wp.cross(a, b)
# compute
r = wp.dot(c, c)
...
可以声明具有不同长度和数据类型的附加向量类型。这在使用 warp.types.vector() 的 Python 作用域 中，在内核之外完成，例如
# declare a new vector type for holding 5 double precision floats:
vec5d = wp.types.vector(length=5, dtype=wp.float64)
声明后，新类型可以在分配数组或在内核内部使用
# create an array of vec5d
arr = wp.zeros(10, dtype=vec5d)
# use inside a kernel
@wp.kernel
def compute( ... ):
# zero initialize a custom named vector type
v = vec5d()
...
# component-wise initialize a named vector type
v = vec5d(wp.float64(1.0),
wp.float64(2.0),
wp.float64(3.0),
wp.float64(4.0),
wp.float64(5.0))
...
此外，可以直接创建这些向量的 匿名 类型实例，而无需提前声明它们的类型。在这种情况下，类型将由构造函数参数推断。例如
@wp.kernel
def compute( ... ):
# zero initialize vector of 5 doubles:
v = wp.vector(dtype=wp.float64, length=5)
# scalar initialize a vector of 5 doubles to the same value:
v = wp.vector(wp.float64(1.0), length=5)
# component-wise initialize a vector of 5 doubles
v = wp.vector(wp.float64(1.0),
wp.float64(2.0),
wp.float64(3.0),
wp.float64(4.0),
wp.float64(5.0))
这些可以与所有标准向量算术运算符一起使用，例如：+、-、标量乘法，并且可以使用具有兼容维度的矩阵进行变换，可能返回不同长度的向量。
矩阵#
还支持任意形状/数值类型的矩阵。内置的具体矩阵类型如下
mat22 mat33 mat44 |
2x2, 3x3, 4x4 单精度浮点矩阵 |
mat22f mat33f mat44f |
2x2, 3x3, 4x4 单精度浮点矩阵 |
mat22d mat33d mat44d |
2x2, 3x3, 4x4 双精度浮点矩阵 |
mat22h mat33h mat44h |
2x2, 3x3, 4x4 半精度浮点矩阵 |
spatial_matrix |
6x6 单精度浮点矩阵 |
spatial_matrixf |
6x6 单精度浮点矩阵 |
spatial_matrixd |
6x6 双精度浮点矩阵 |
spatial_matrixh |
6x6 半精度浮点矩阵 |
矩阵以行主序格式存储，并支持大多数标准线性代数运算
@wp.kernel
def compute( ... ):
# initialize matrix
m = wp.mat22(1.0, 2.0,
3.0, 4.0)
# compute inverse
minv = wp.inverse(m)
# transform vector
v = minv * wp.vec2(0.5, 0.3)
...
与向量类似，可以使用 wp.types.matrix() 声明具有任意形状和数据类型的新矩阵类型，例如
# declare a new 3x2 half precision float matrix type:
mat32h = wp.types.matrix(shape=(3,2), dtype=wp.float64)
# create an array of this type
a = wp.zeros(10, dtype=mat32h)
这些可以在内核内部使用
@wp.kernel
def compute( ... ):
...
# initialize a mat32h matrix
m = mat32h(wp.float16(1.0), wp.float16(2.0),
wp.float16(3.0), wp.float16(4.0),
wp.float16(5.0), wp.float16(6.0))
# declare a 2 component half precision vector
v2 = wp.vec2h(wp.float16(1.0), wp.float16(1.0))
# multiply by the matrix, returning a 3 component vector:
v3 = m * v2
...
还可以在内核内部直接创建匿名类型的实例，其中类型由构造函数参数推断，如下所示
@wp.kernel
def compute( ... ):
...
# create a 3x2 half precision matrix from components (row major ordering):
m = wp.matrix(
wp.float16(1.0), wp.float16(2.0),
wp.float16(1.0), wp.float16(2.0),
wp.float16(1.0), wp.float16(2.0),
shape=(3,2))
# zero initialize a 3x2 half precision matrix:
m = wp.matrix(wp.float16(0.0),shape=(3,2))
# create a 5x5 double precision identity matrix:
m = wp.identity(n=5, dtype=wp.float64)
与向量一样，你可以使用这些变量进行标准的矩阵算术运算，以及将具有兼容形状的矩阵相乘，并可能返回具有新形状的矩阵。
四元数#
Warp 支持布局为 i, j, k, w 的四元数，其中 w 是实部。以下是内置的具体四元数类型
quat |
单精度浮点四元数 |
quatf |
单精度浮点四元数 |
quatd |
双精度浮点四元数 |
quath |
半精度浮点四元数 |
四元数可以用来变换向量，如下所示
@wp.kernel
def compute( ... ):
...
# construct a 30 degree rotation around the x-axis
q = wp.quat_from_axis_angle(wp.vec3(1.0, 0.0, 0.0), wp.degrees(30.0))
# rotate an axis by this quaternion
v = wp.quat_rotate(q, wp.vec3(0.0, 1.0, 0.0))
与向量和矩阵一样，你可以像这样声明具有任意数值类型的四元数类型
quatd = wp.types.quaternion(dtype=wp.float64)
你也可以在内核内部创建单位四元数和匿名类型的实例，如下所示
@wp.kernel
def compute( ... ):
...
# create a double precision identity quaternion:
qd = wp.quat_identity(dtype=wp.float64)
# precision defaults to wp.float32 so this creates a single precision identity quaternion:
qf = wp.quat_identity()
# create a half precision quaternion from components, or a vector/scalar:
qh = wp.quaternion(wp.float16(0.0),
wp.float16(0.0),
wp.float16(0.0),
wp.float16(1.0))
qh = wp.quaternion(
wp.vector(wp.float16(0.0),wp.float16(0.0),wp.float16(0.0)),
wp.float16(1.0))
变换#
变换是表示空间刚体变换的 7D 浮点向量，格式为 (p, q)，其中 p 是 3D 向量，q 是四元数。
transform |
单精度浮点变换 |
transformf |
单精度浮点变换 |
transformd |
双精度浮点变换 |
transformh |
半精度浮点变换 |
变换可以在内核内部由平移和旋转部分构造
@wp.kernel
def compute( ... ):
...
# create a transform from a vector/quaternion:
t = wp.transform(
wp.vec3(1.0, 2.0, 3.0),
wp.quat_from_axis_angle(wp.vec3(0.0, 1.0, 0.0), wp.degrees(30.0)))
# transform a point
p = wp.transform_point(t, wp.vec3(10.0, 0.5, 1.0))
# transform a vector (ignore translation)
p = wp.transform_vector(t, wp.vec3(10.0, 0.5, 1.0))
与向量和矩阵一样，你可以使用 wp.types.transformation() 声明具有任意数值类型的变换类型，例如
transformd = wp.types.transformation(dtype=wp.float64)
你也可以在内核内部创建单位变换和匿名类型的实例，如下所示
@wp.kernel
def compute( ... ):
# create double precision identity transform:
qd = wp.transform_identity(dtype=wp.float64)
结构体#
用户可以使用 @wp.struct 装饰器定义自定义结构体类型，如下所示
@wp.struct
class MyStruct:
param1: int
param2: float
param3: wp.array(dtype=wp.vec3)
结构体属性必须用其相应的类型进行标注。它们可以在 Python 作用域中构造，然后作为参数传递给内核
@wp.kernel
def compute(args: MyStruct):
tid = wp.tid()
print(args.param1)
print(args.param2)
print(args.param3[tid])
# construct an instance of the struct in Python
s = MyStruct()
s.param1 = 10
s.param2 = 2.5
s.param3 = wp.zeros(shape=10, dtype=wp.vec3)
# pass to our compute kernel
wp.launch(compute, dim=10, inputs=[s])
结构体数组可以按如下方式进行零初始化
a = wp.zeros(shape=10, dtype=MyStruct)
结构体数组也可以通过结构体对象列表进行初始化
a = wp.array([MyStruct(), MyStruct(), MyStruct()], dtype=MyStruct)
示例：在梯度计算中使用结构体#
import numpy as np
import warp as wp
@wp.struct
class TestStruct:
x: wp.vec3
a: wp.array(dtype=wp.vec3)
b: wp.array(dtype=wp.vec3)
@wp.kernel
def test_kernel(s: TestStruct):
tid = wp.tid()
s.b[tid] = s.a[tid] + s.x
@wp.kernel
def loss_kernel(s: TestStruct, loss: wp.array(dtype=float)):
tid = wp.tid()
v = s.b[tid]
wp.atomic_add(loss, 0, float(tid + 1) * (v[0] + 2.0 * v[1] + 3.0 * v[2]))
# create struct
ts = TestStruct()
# set members
ts.x = wp.vec3(1.0, 2.0, 3.0)
ts.a = wp.array(np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), dtype=wp.vec3, requires_grad=True)
ts.b = wp.zeros(2, dtype=wp.vec3, requires_grad=True)
loss = wp.zeros(1, dtype=float, requires_grad=True)
tape = wp.Tape()
with tape:
wp.launch(test_kernel, dim=2, inputs=[ts])
wp.launch(loss_kernel, dim=2, inputs=[ts, loss])
tape.backward(loss)
print(loss)
print(ts.a)
[120.]
[[1. 2. 3.]
[4. 5. 6.]]
示例：定义运算符重载#
@wp.struct
class Complex:
real: float
imag: float
@wp.func
def add(
a: Complex,
b: Complex,
) -> Complex:
return Complex(a.real + b.real, a.imag + b.imag)
@wp.func
def mul(
a: Complex,
b: Complex,
) -> Complex:
return Complex(
a.real * b.real - a.imag * b.imag,
a.real * b.imag + a.imag * b.real,
)
@wp.kernel
def kernel():
a = Complex(1.0, 2.0)
b = Complex(3.0, 4.0)
c = a + b
wp.printf("%.0f %+.0fi\n", c.real, c.imag)
d = a * b
wp.printf("%.0f %+.0fi\n", d.real, d.imag)
wp.launch(kernel, dim=(1,))
wp.synchronize()
类型转换#
Warp 对类型转换特别严格，并且不执行数值类型之间的 任何 隐式转换。用户负责确保大多数算术运算符的类型匹配，例如：x = float(0.0) + int(4) 将导致错误。对于习惯 C 风格转换的用户来说，这可能令人惊讶，但这避免了一类由隐式转换引起的常见错误。
用户应使用 int()、float()、wp.float16()、wp.uint8() 等构造函数将变量显式转换为兼容类型。
注意
出于性能原因，Warp 依赖于原生编译器来执行数值转换（例如，CPU 使用 LLVM，CUDA 使用 NVRTC）。这通常不是问题，但在某些情况下，结果可能因设备而异。例如，转换 wp.uint8(-1.0) 会导致未定义行为，因为浮点值 -1.0 超出了无符号整型的范围。
C++ 编译器可以自由地按其认为合适的方式处理此类情况。只有当被转换的值在目标数据类型支持的范围内时，才能保证数值转换产生正确的结果。
常量#
Warp 内核可以访问在内核外部定义的 Python 变量，这些变量在内核内部被视为编译时常量。
TYPE_SPHERE = wp.constant(0)
TYPE_CUBE = wp.constant(1)
TYPE_CAPSULE = wp.constant(2)
@wp.kernel
def collide(geometry: wp.array(dtype=int)):
t = geometry[wp.tid()]
if t == TYPE_SPHERE:
print("sphere")
elif t == TYPE_CUBE:
print("cube")
elif t == TYPE_CAPSULE:
print("capsule")
请注意，不再需要使用 wp.constant()，但它会执行一些类型检查，并且可以提醒这些变量应作为 Warp 常量使用。
当引用的 Python 变量不变时，其行为简单直观。有关详细信息和更复杂的场景，请参阅 外部引用和常量。代码生成 部分包含高级用法的附加信息和技巧。
预定义常量#
为了方便起见，Warp 提供了许多预定义的数学常量，这些常量可以在 Warp 内核内部和外部使用。下表中的常量也定义了小写版本，例如 wp.E 和 wp.e 是等效的。
名称 |
值 |
|---|---|
wp.E |
2.71828182845904523536 |
wp.LOG2E |
1.44269504088896340736 |
wp.LOG10E |
0.43429448190325182765 |
wp.LN2 |
0.69314718055994530942 |
wp.LN10 |
2.30258509299404568402 |
wp.PHI |
1.61803398874989484820 |
wp.PI |
3.14159265358979323846 |
wp.HALF_PI |
1.57079632679489661923 |
wp.TAU |
6.28318530717958647692 |
wp.INF |
math.inf |
wp.NAN |
float(‘nan’) |
wp.NAN 常量只能用于浮点类型。涉及 wp.NAN 的比较遵循 IEEE 754 标准，例如 wp.float32(wp.NAN) == wp.float32(wp.NAN) 返回 False。内置函数 wp.isnan() 可用于确定一个值是否为 NaN（或者向量、矩阵或四元数是否包含 NaN 条目）。
以下示例展示了如何使用 wp.inf 常量在 Warp 中使用正无穷和负无穷与浮点类型。
@wp.kernel
def test_infinity(outputs: wp.array(dtype=wp.float32)):
outputs[0] = wp.float32(wp.inf) # inf
outputs[1] = wp.float32(-wp.inf) # -inf
outputs[2] = wp.float32(2.0 * wp.inf) # inf
outputs[3] = wp.float32(-2.0 * wp.inf) # -inf
outputs[4] = wp.float32(2.0 / 0.0) # inf
outputs[5] = wp.float32(-2.0 / 0.0) # -inf
运算符#
布尔运算符#
a and b |
如果 a 和 b 都为 True 则为 True |
a or b |
如果 a 或 b 为 True 则为 True |
not a |
如果 a 为 False 则为 True，否则为 False |
注意
像 if (a and b): 这样的表达式目前不执行短路求值。在这种情况下，即使 a 为 False，b 也会被求值。用户应注意确保在所有情况下次要条件都是安全的（例如：不要索引越界）。
比较运算符#
a > b |
如果 a 严格大于 b 则为 True |
a < b |
如果 a 严格小于 b 则为 True |
a >= b |
如果 a 大于或等于 b 则为 True |
a <= b |
如果 a 小于或等于 b 则为 True |
a == b |
如果 a 等于 b 则为 True |
a != b |
如果 a 不等于 b 则为 True |
算术运算符#
a + b |
加法 |
a - b |
减法 |
a * b |
乘法 |
a @ b |
矩阵乘法 |
a / b |
浮点除法 |
a // b |
向下取整除法 |
a ** b |
幂运算 |
a % b |
求余 |
由于 Warp 不执行隐式类型转换，操作数应具有兼容的数据类型。用户应使用 float()、int()、wp.int64() 等类型构造函数将变量转换为正确的类型。
乘法表达式 a * b 也可用于在 矩阵类型 之间执行矩阵乘法。
流#
CUDA 流是在 GPU 上按顺序执行的一系列操作。来自不同流的操作可以并发运行，并可能由设备调度程序交错执行。有关使用流的更多信息，请参阅 流文档。
- class warp.Stream(*args, **kwargs)[源]#
- __init__(device=None, priority=0, **kwargs)[源]#
在设备上初始化流，可选择指定优先级。
- 参数：
- Raises:
RuntimeError – 如果在 Warp 完成初始化之前调用函数，并且
device不是Device的实例。RuntimeError –
device不是 CUDA 设备。RuntimeError – 无法在设备上创建流。
TypeError – 请求的流优先级不是整数。
- record_event(event=None)[源]#
在流上记录一个事件。
- 参数：
event (Event | None) – 要记录到流上的 warp.Event 实例。如果未提供，则将在同一设备上创建一个
Event实例。- Raises:
RuntimeError – 提供的
Event来自与记录流不同的设备。- 返回类型：
- warp.get_stream(device=None)[源]#
返回指定设备当前使用的流。
- 参数：
device (Device | str | None) – 一个可选的
Device实例或设备别名（例如“cuda:0”），将返回其当前流。如果为None，将使用默认设备。- Raises:
RuntimeError – 设备不是 CUDA 设备。
- 返回类型：
- warp.wait_stream(other_stream, event=None)[源]#
调用当前流上的
Stream.wait_stream()的便捷函数。
- warp.synchronize_stream(stream_or_device=None)[源]#
将调用 CPU 线程与指定流上任何未完成的 CUDA 工作同步。
此函数允许主机应用程序代码确保流上的所有内核启动和内存复制都已完成。
事件#
事件可以插入到流中，并用于将一个流与另一个流同步。有关如何使用事件进行跨流同步的信息，请参阅 事件文档，有关如何使用事件测量 GPU 性能的信息，请参阅 CUDA 事件计时文档。
- class warp.Event(*args, **kwargs)[源]#
一个可以记录到流上的 CUDA 事件。
事件可用于设备端同步，这不会阻塞主机线程。
- __init__(
- device=None,
- cuda_event=None,
- enable_timing=False,
- interprocess=False,
在 CUDA 设备上初始化事件。
- 参数：
device (Device | str | None) – CUDA 设备，该事件可以记录到其流中。如果为
None，则将使用当前默认设备。cuda_event – 指向先前分配的 CUDA 事件的指针。如果为 None，则将在关联的设备上分配新事件。
enable_timing (bool) – 如果为
True，此事件将记录计时数据。get_event_elapsed_time()可用于测量两个以enable_timing=True创建并记录到流上的事件之间的时间。interprocess (bool) – 如果为
True，此事件可用作进程间事件。
- Raises:
RuntimeError – 事件无法创建。
ValueError – 不允许
enable_timing=True和interprocess=True的组合。
- ipc_handle()[source]#
返回事件的 CUDA IPC 句柄，作为 64 字节的
bytes对象。事件必须已使用
interprocess=True创建，才能获得有效的进程间句柄。IPC 目前仅在 Linux 上支持。
示例
创建事件并获取其 IPC 句柄
e1 = wp.Event(interprocess=True) event_handle = e1.ipc_handle()
- Raises:
RuntimeError – 设备不支持 IPC。
- 返回类型：
- warp.record_event(event=None)[source]#
在当前流上调用
Stream.record_event()的便利函数。
- warp.wait_event(event)[source]#
在当前流上调用
Stream.wait_event()的便利函数。
- warp.synchronize_event(event)[source]#
将调用 CPU 线程与记录在 CUDA 流上的事件同步。
此函数允许主机应用程序代码确保已达到特定的同步点。
- 参数：
event (Event) – 要等待的事件。
- warp.get_event_elapsed_time(start_event, end_event, synchronize=True)[source]#
获取两个已记录事件之间的流逝时间。
两个事件都必须已通过
record_event()或warp.Stream.record_event()记录。如果
synchronize为 False，则调用者必须确保在调用get_event_elapsed_time()之前设备执行已到达end_event。
图 (Graphs)#
与 C++ 或原生程序相比，从 Python 启动 kernels 会引入显著的额外开销。为了解决这个问题，Warp 提供了 CUDA graphs 的概念，允许记录大批量 kernels 并以极低的 CPU 开销重播它们。
要记录一系列 kernel 启动，请按如下方式使用 wp.capture_begin() 和 wp.capture_end() API
# begin capture
wp.capture_begin(device="cuda")
try:
# record launches
for i in range(100):
wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
finally:
# end capture and return a graph object
graph = wp.capture_end(device="cuda")
我们强烈建议在捕获图时使用 try-finally 模式，因为 finally 语句将确保调用 wp.capture_end，即使在捕获期间发生异常，否则会使流陷入捕获状态。
构建图后即可执行
wp.capture_launch(graph)
wp.ScopedCapture 上下文管理器可用于简化代码并确保无论是否发生异常都会调用 wp.capture_end
with wp.ScopedCapture(device="cuda") as capture:
# record launches
for i in range(100):
wp.launch(kernel=compute1, inputs=[a, b], device="cuda")
wp.capture_launch(capture.graph)
请注意，图中仅记录启动调用；在 kernel 代码之外执行的任何 Python 代码都不会被记录。通常，只有当图将被重用或多次启动时，使用 CUDA graphs 才是有益的，因为存在图创建的开销。
- warp.capture_begin(
- device=None,
- stream=None,
- force_module_load=None,
- external=False,
开始捕获 CUDA 图
捕获 CUDA 设备上的所有后续 kernel 启动和内存操作。这可用于记录大量 kernels 并以低开销重播它们。
如果指定了 device，捕获将开始于当前与该设备关联的 CUDA 流。如果指定了 stream，捕获将开始于给定的流。如果两者都省略，捕获将开始于当前设备的当前流。
- 参数：
stream (Stream | None) – 要在其上捕获的 CUDA 流
force_module_load (bool | None) – 在捕获前是否强制加载所有 kernels。一般来说，最好使用
load_module()选择性加载 kernels。当使用支持 CUDA 12.3 或更新版本的 CUDA 驱动程序运行时，不建议将此选项设置为True，因为在更新的驱动程序上可以在图捕获期间加载 kernels。如果此参数为None，则如果驱动程序早于 CUDA 12.3，其行为将继承自wp.config.enable_graph_capture_module_load_by_default。external (bool) – 捕获是否已由外部开始
- warp.capture_end(device=None, stream=None)[source]#
结束 CUDA 图的捕获。
- 参数：
- 返回值：
一个
Graph对象，可使用capture_launch()启动- 返回类型：
图 (Graph)
- warp.capture_launch(graph, stream=None)[source]#
启动先前捕获的 CUDA 图
- 参数：
graph (Graph) – 由
capture_end()返回的Graph
网格 (Meshes)#
Warp 提供了一个 wp.Mesh 类来管理三角形网格数据。要创建网格，用户需要提供点、索引以及可选的速度数组
mesh = wp.Mesh(points, indices, velocities)
注意
Mesh 对象维护对其输入几何体缓冲区的引用。所有缓冲区应驻留在同一设备上。
可以通过其 id 属性将网格传递给 kernels，该属性是一个唯一标识网格的 uint64 值。在 kernel 内部，您可以对网格执行几何查询，例如光线投射或最近点查找
@wp.kernel
def raycast(mesh: wp.uint64,
ray_origin: wp.array(dtype=wp.vec3),
ray_dir: wp.array(dtype=wp.vec3),
ray_hit: wp.array(dtype=wp.vec3)):
tid = wp.tid()
t = float(0.0) # hit distance along ray
u = float(0.0) # hit face barycentric u
v = float(0.0) # hit face barycentric v
sign = float(0.0) # hit face sign
n = wp.vec3() # hit face normal
f = int(0) # hit face index
color = wp.vec3()
# ray cast against the mesh
if wp.mesh_query_ray(mesh, ray_origin[tid], ray_dir[tid], 1.e+6, t, u, v, sign, n, f):
# if we got a hit then set color to the face normal
color = n*0.5 + wp.vec3(0.5, 0.5, 0.5)
ray_hit[tid] = color
用户可以通过简单地修改点缓冲区来在运行时更新网格顶点位置。修改点位置后，用户应调用 Mesh.refit() 以重建边界体积层次结构 (BVH)，确保查询正常工作。
注意
目前不支持在运行时更新网格拓扑（索引）。用户应改为重新创建一个新的 Mesh 对象。
- class warp.Mesh(*args, **kwargs)[source]#
- __init__(
- points,
- indices,
- velocities=None,
- support_winding_number=False,
- bvh_constructor=None,
表示三角形网格的类。
- id#
此网格对象的唯一标识符，可以传递给 kernels。
- device#
此对象所在的设备，所有缓冲区必须驻留在同一设备上。
- 参数：
points (array) – 数据类型为
warp.vec3的顶点位置数组。indices (array) – 数据类型为
warp.int32的三角形索引数组。应为一个形状为(num_tris * 3)的 1D 数组。velocities (array | None) – 可选的顶点速度数组，数据类型为
warp.vec3。support_winding_number (bool) – 如果为
True，网格将构建额外的数据结构以支持wp.mesh_query_point_sign_winding_number()查询。bvh_constructor (str | None) – 底层 BVH 的构造算法（参见
Bvh的文档字符串了解解释）。有效选择为"sah"、"median"、"lbvh"或None。
哈希网格 (Hash Grids)#
许多基于粒子的模拟方法，例如离散元法 (DEM) 或平滑粒子流体动力学 (SPH)，都涉及迭代空间邻居来计算力相互作用。哈希网格是一种成熟的数据结构，可加速这些最近邻居查询，尤其适用于 GPU。
为了支持空间邻居查询，Warp 提供了一个 HashGrid 对象，可以按如下方式创建
grid = wp.HashGrid(dim_x=128, dim_y=128, dim_z=128, device="cuda")
grid.build(points=p, radius=r)
p 是一个 wp.vec3 点位置数组，r 是用于构建网格的半径。然后可以使用 wp.hash_grid_query() 和 wp.hash_grid_query_next() 在 kernel 代码内部迭代邻居，如下所示
@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
注意
HashGrid 查询将返回查询半径内所有单元格中的点。当存在哈希冲突时，意味着将返回一些查询半径之外的点，用户应在 kernel 内部自行检查距离。查询不自行检查每个返回点距离的原因是，kernels 通常会自行计算距离，因此两次检查/计算距离会是冗余的。
- class warp.HashGrid(*args, **kwargs)[source]#
- __init__(dim_x, dim_y, dim_z, device=None)[source]#
表示用于加速点查询的哈希网格对象的类。
- id#
此网格对象的唯一标识符，可以传递给 kernels。
- device#
此对象所在的设备，所有缓冲区必须驻留在同一设备上。
- build(points, radius)[source]#
更新哈希网格数据结构。
此方法重建底层数据结构，应在点集发生任何更改时调用。
- 参数：
points (
warp.array) – 类型为warp.vec3的点数组radius (float) – 用于对点进行分桶的单元格大小，单元格是边长为此宽度的立方体。为了获得最佳性能，用于构造网格的半径应与执行查询时使用的半径紧密匹配。
体素 (Volumes)#
稀疏体素对于在大域上表示网格数据非常有用，例如复杂物体的符号距离场 (SDF) 或大规模流体流动的速度。Warp 支持读取使用 NanoVDB 标准存储的稀疏体素网格。用户可以直接访问体素，或使用内置的最近点或三线性插值从世界或局部空间采样网格数据。
可以从包含 NanoVDB 网格的 Warp 数组、使用 load_from_nvdb() 从标准 .nvdb 文件内容、使用 load_from_address() 从未压缩的内存缓冲区，或使用 load_from_numpy() 从密集的 3D NumPy 数组直接创建 Volume 对象。
还可以使用 allocate()、allocate_by_tiles() 或 allocate_by_voxels() 创建 Volume 对象。可以使用 wp.volume_store() 在 Warp kernel 中修改 Volume 对象的值。
注意
Warp 目前不支持在运行时修改稀疏体素的拓扑结构。
下面我们给出一个从现有 NanoVDB 文件创建 Volume 对象的示例
# open NanoVDB file on disk
file = open("mygrid.nvdb", "rb")
# create Volume object
volume = wp.Volume.load_from_nvdb(file, device="cpu")
注意
NanoVDB 库编写的文件（通常标记为 .nvdb 扩展名）可以包含具有各种压缩方法的多个网格，但一个 Volume 对象表示单个 NanoVDB 网格。默认加载第一个网格，然后可以通过重复调用 load_next_grid() 创建文件中对应于其他网格的 Warp volumes。NanoVDB 的未压缩和 zip 压缩文件格式开箱即用，blosc 压缩文件需要安装 blosc Python 包。
要在 kernel 内部采样体素，我们通过 ID 将其引用传递给它，并使用内置采样模式
@wp.kernel
def sample_grid(volume: wp.uint64,
points: wp.array(dtype=wp.vec3),
samples: wp.array(dtype=float)):
tid = wp.tid()
# load sample point in world-space
p = points[tid]
# transform position to the volume's local-space
q = wp.volume_world_to_index(volume, p)
# sample volume with trilinear interpolation
f = wp.volume_sample(volume, q, wp.Volume.LINEAR, dtype=float)
# write result
samples[tid] = f
Warp 还支持 NanoVDB 索引网格，它们提供了一种内存高效的体素索引线性化，可以引用任意形状数组中的值
@wp.kernel
def sample_index_grid(volume: wp.uint64,
points: wp.array(dtype=wp.vec3),
voxel_values: wp.array(dtype=Any)):
tid = wp.tid()
# load sample point in world-space
p = points[tid]
# transform position to the volume's local-space
q = wp.volume_world_to_index(volume, p)
# sample volume with trilinear interpolation
background_value = voxel_values.dtype(0.0)
f = wp.volume_sample_index(volume, q, wp.Volume.LINEAR, voxel_values, background_value)
可以使用 get_voxels() 恢复所有可索引体素的坐标。NanoVDB 网格还可能包含嵌入的盲数据数组；可以使用 feature_array() 函数访问这些数组。
- class warp.Volume(*args, **kwargs)[source]#
- CLOSEST = 0#
指定采样期间最近邻插值的枚举值
- LINEAR = 1#
在采样期间指定三线性插值的枚举值
- class GridInfo(
- name,
- size_in_bytes,
- grid_index,
- grid_count,
- type_str,
- translation,
- transform_matrix,
网格元数据
- 参数：
- translation: vec3f#
索引到世界的平移
- transform_matrix: mat33f#
索引到世界变换的线性部分
- property dtype: type[source]#
体素值的类型，作为 Warp 类型。
如果网格不包含值（例如，索引网格），或者 NanoVDB 类型无法表示为 Warp 类型，则返回
None。
- get_feature_array_info(feature_index)[source]#
返回与
feature_index处的特征数组关联的元数据。- 参数：
feature_index (int)
- 返回类型：
- classmethod load_from_nvdb(file_or_buffer, device=None)[source]#
从序列化的 NanoVDB 文件或内存缓冲区创建
Volume对象。- 返回值：
一个
warp.Volume对象。- 返回类型：
- save_to_nvdb(path, codec='none')[source]#
将体素序列化为 NanoVDB (.nvdb) 文件。
- 参数：
path – 要保存的文件路径。
codec (Literal['none', 'zip', 'blosc']) – 使用的压缩编码器，“none” - 无压缩，“zip” - ZIP 压缩，“blosc” - BLOSC 压缩，需要安装 blosc 模块。
- classmethod load_from_address(grid_ptr, buffer_size=0, device=None)[source]#
创建一个新的
Volume，别名内存中的网格缓冲区。与用于加载序列化 NanoVDB 网格的
load_from_nvdb()不同，此处缓冲区必须是未压缩的，并且不得包含文件头信息。如果传递的地址不包含 NanoVDB 网格，则此函数的行为未定义。- 参数：
- 返回类型：
返回新创建的体素。
- load_next_grid()[source]#
尝试为与此体素链接的下一个网格创建一个新的 warp 体素。
下一个网格是否存在，是根据 grid_index 和 grid_count 元数据以及此体素在内存中缓冲区的大小推断的。
返回新创建的体素，如果不存在下一个网格，则返回 None。
- 返回类型：
- classmethod load_from_numpy(
- ndarray,
- min_world=(0.0, 0.0, 0.0),
- voxel_size=1.0,
- bg_value=0.0,
- device=None,
从密集 3D NumPy 数组创建体素对象。
此函数仅支持 CUDA 设备。
- 参数：
min_world – 体素下角的 3D 坐标。
voxel_size – 每个体素在空间坐标中的大小。
bg_value – 背景值
device – 创建体素的 CUDA 设备，例如：“cuda” 或 “cuda:0”。
ndarray (array)
- 返回值：
一个
warp.Volume对象。- 返回类型：
- classmethod allocate(
- min,
- max,
- voxel_size,
- bg_value=0.0,
- translation=(0.0, 0.0, 0.0),
- points_in_world_space=False,
- device=None,
根据 min 和 max 定义的边界框分配新的体素。
此函数仅支持 CUDA 设备。
分配一个足够大的体素以包含体素 [min[0], min[1], min[2]] - [max[0], max[1], max[2]]（包括边界）。如果 points_in_world_space 为 true，则首先使用给定的体素大小和平移将 min 和 max 转换为索引空间，然后分配具有这些坐标的体素。
最小分配单位是 8x8x8 体素的密集瓦片，请求的边界框向上取整到瓦片，并且生成的瓦片将在新的体素中可用。
- 参数：
min (array-like) – 边界框在索引空间或世界空间中的下部 3D 坐标，包括边界。
max (array-like) – 边界框在索引空间或世界空间中的上部 3D 坐标，包括边界。
voxel_size (float) – 新体素的体素大小。
bg_value (float or array-like) – 体素未分配体素的值，也定义体素的类型。如果此值为 array-like，则创建
warp.vec3体素，否则创建 float 体素。translation (array-like) – 索引空间和世界空间之间的平移。
device (Devicelike) – 创建体素的 CUDA 设备，例如：“cuda” 或 “cuda:0”。
- 返回类型：
- classmethod allocate_by_tiles(
- tile_points,
- voxel_size=None,
- bg_value=0.0,
- translation=(0.0, 0.0, 0.0),
- device=None,
- transform=None,
为每个点 tile_points 分配具有活动瓦片的新体素。
此函数仅支持 CUDA 设备。
最小分配单位是 8x8x8 体素的密集瓦片。这是分配稀疏体素的主要方法。它使用一个点数组来指示必须分配的瓦片。
- 示例用例
tile_points 可以直接在索引空间中标记瓦片，就像此方法由 allocate 调用时一样。
tile_points 可以是模拟中使用的点列表，需要将数据传输到体素。
- 参数：
tile_points (
warp.array) – 定义要分配的瓦片的点数组。数组可以使用整数标量类型（warp.int32的 2D N×3 数组或 warp.vec3i 值的一维数组），指示索引空间位置，或浮点标量类型（warp.float32的 2D N×3 数组或 warp.vec3f 值的一维数组），指示世界空间位置。允许每个瓦片存在重复的点，并将高效去重。voxel_size (float or array-like) – 新体素的体素大小。如果给出 transform 则忽略。
bg_value (array-like, scalar or None) – 体素未分配体素的值，也定义体素的类型。如果 bg_value 是
None，则创建索引体素。其他支持的网格类型包括 int、float、vec3f 和 vec4f。translation (array-like) – 索引空间和世界空间之间的平移。
transform (array-like) – 索引空间和世界空间之间的线性变换。如果为
None，则从 voxel_size 推导。device (Devicelike) – 创建体素的 CUDA 设备，例如：“cuda” 或 “cuda:0”。
- 返回类型：
- classmethod allocate_by_voxels(
- voxel_points,
- voxel_size=None,
- translation=(0.0, 0.0, 0.0),
- device=None,
- transform=None,
为每个点 voxel_points 分配具有活动体素的新体素。
此函数创建一个索引体素，这是一种特殊的体素，它不存储任何显式负载，而是为每个活动体素编码一个线性化索引，允许从任意外部数组查找和采样数据。
此函数仅支持 CUDA 设备。
- 参数：
voxel_points (
warp.array) – 定义要分配的体素的点数组。数组可以使用整数标量类型（warp.int32的 2D N×3 数组或 warp.vec3i 值的一维数组），指示索引空间位置，或浮点标量类型（warp.float32的 2D N×3 数组或 warp.vec3f 值的一维数组），指示世界空间位置。允许每个瓦片存在重复的点，并将高效去重。voxel_size (float or array-like) – 新体素的体素大小。如果给出 transform 则忽略。
translation (array-like) – 索引空间和世界空间之间的平移。
transform (array-like) – 索引空间和世界空间之间的线性变换。如果为
None，则从 voxel_size 推导。device (Devicelike) – 创建体素的 CUDA 设备，例如：“cuda” 或 “cuda:0”。
- 返回类型：
另请参阅
内核中可用的体素函数的参考。
边界值层级结构 (BVH)#
wp.Bvh 类可用于为一组边界体创建 BVH。然后可以遍历此对象，使用 bvh_query_ray() 确定哪些部分与射线相交，并使用 bvh_query_aabb() 确定哪些部分与某个边界体素重叠。
以下代码片段演示了如何从 100 个随机边界体素创建 wp.Bvh 对象
rng = np.random.default_rng(123)
num_bounds = 100
lowers = rng.random(size=(num_bounds, 3)) * 5.0
uppers = lowers + rng.random(size=(num_bounds, 3)) * 5.0
device_lowers = wp.array(lowers, dtype=wp.vec3, device="cuda:0")
device_uppers = wp.array(uppers, dtype=wp.vec3, device="cuda:0")
bvh = wp.Bvh(device_lowers, device_uppers)
- class warp.Bvh(*args, **kwargs)[source]#
- __init__(lowers, uppers, constructor=None)[source]#
表示边界值层级结构的类。
根据输入边界所在的设备，它可以是 CPU 树或 GPU 树。
- id#
此 BVH 对象的唯一标识符，可传递给内核。
- device#
此对象所在的设备，所有缓冲区必须驻留在同一设备上。
- 参数：
注意
BVH 构造函数说明
"sah": 一种基于 CPU 的自顶向下构造函数，其中 AABB 基于表面积启发式 (SAH) 进行分割。构建时间比其他方法稍长，但查询性能最佳。"median": 一种基于 CPU 的自顶向下构造函数，其中 AABB 基于 AABB 中图元的质心中位数进行分割。此构造函数比 SAH 快，但查询性能较差。"lbvh": 一种基于 GPU 的自底向上构造函数，可最大限度地提高并行性。构建速度非常快，尤其对于大型模型。查询性能比"sah"稍慢。None: 构造函数将根据树所在的设备自动选择。对于 GPU 树，将选择"lbvh"构造函数；对于 CPU 树，将选择"sah"构造函数。
GPU 树支持所有三种构造函数。当为 GPU 树选择基于 CPU 的构造函数时，边界将被复制回 CPU 以运行基于 CPU 的构造函数。构建完成后，CPU 树将被复制到 GPU。
CPU 树仅支持
"sah"和"median"。如果为 CPU 树选择了"lbvh"，将发出警告消息，并且构造函数将自动回退到"sah"。
示例：BVH 光线遍历#
对数据结构执行光线遍历的示例如下
@wp.kernel
def bvh_query_ray(
bvh_id: wp.uint64,
start: wp.vec3,
dir: wp.vec3,
bounds_intersected: wp.array(dtype=wp.bool),
):
query = wp.bvh_query_ray(bvh_id, start, dir)
bounds_nr = wp.int32(0)
while wp.bvh_query_next(query, bounds_nr):
# The ray intersects the volume with index bounds_nr
bounds_intersected[bounds_nr] = True
bounds_intersected = wp.zeros(shape=(num_bounds), dtype=wp.bool, device="cuda:0")
query_start = wp.vec3(0.0, 0.0, 0.0)
query_dir = wp.normalize(wp.vec3(1.0, 1.0, 1.0))
wp.launch(
kernel=bvh_query_ray,
dim=1,
inputs=[bvh.id, query_start, query_dir, bounds_intersected],
device="cuda:0",
)
Warp 核函数 bvh_query_ray 使用单个线程启动，提供了 wp.Bvh 对象的唯一 uint64 标识符、描述光线的参数以及一个存储结果的数组。在 bvh_query_ray 中，wp.bvh_query_ray() 被调用一次以获取存储在变量 query 中的对象。还分配了一个整数作为 bounds_nr 来存储遍历的卷索引。一个 while 语句用于使用 wp.bvh_query_next() 执行实际遍历，只要存在相交的边界，该函数就返回 True。
示例：BVH 体积遍历#
与光线遍历示例类似，我们可以执行体积遍历以查找完全包含在指定边界框内的体积。
@wp.kernel
def bvh_query_aabb(
bvh_id: wp.uint64,
lower: wp.vec3,
upper: wp.vec3,
bounds_intersected: wp.array(dtype=wp.bool),
):
query = wp.bvh_query_aabb(bvh_id, lower, upper)
bounds_nr = wp.int32(0)
while wp.bvh_query_next(query, bounds_nr):
# The volume with index bounds_nr is fully contained
# in the (lower,upper) bounding box
bounds_intersected[bounds_nr] = True
bounds_intersected = wp.zeros(shape=(num_bounds), dtype=wp.bool, device="cuda:0")
query_lower = wp.vec3(4.0, 4.0, 4.0)
query_upper = wp.vec3(6.0, 6.0, 6.0)
wp.launch(
kernel=bvh_query_aabb,
dim=1,
inputs=[bvh.id, query_lower, query_upper, bounds_intersected],
device="cuda:0",
)
核函数与光线遍历示例几乎相同，只是我们使用 wp.bvh_query_aabb() 获取 query。
Marching Cubes（行进立方）#
wp.MarchingCubes 类可用于提取近似 3D 标量场等值面的 2D 网格。可以使用 warp.renderer.UsdRenderer 将生成的三角形网格保存到 USD 文件。
有关用法示例，请参见 warp/examples/core/example_marching_cubes.py。
- class warp.MarchingCubes(*args, **kwargs)[source]#
- __init__(nx, ny, nz, max_verts, max_tris, device=None)[source]#
基于 CUDA 的行进立方算法，用于从 3D 体积中提取 2D 表面网格。
- id#
此对象的唯一标识符。
- verts#
输出表面网格的顶点位置数组，类型为
warp.vec3f。运行surface()后会填充此数组。- 类型：
- indices#
包含类型为
warp.int32的索引数组，用于定义输出表面网格的三角形。运行surface()后会填充此数组。数组中每三个连续的整数代表一个三角形，其中每个整数都是
verts数组中顶点的索引。- 类型：
- 参数：
- Raises:
RuntimeError –
device不是 CUDA 设备。
注意
行进立方的形状应与要生成表面的标量场的形状匹配。
- resize(nx, ny, nz, max_verts, max_tris)[source]#
更新行进立方计算的预期输入和最大输出大小。
此函数对底层缓冲区没有立即影响。新值在下次调用
surface()时生效。
- surface(field, threshold)[source]#
从 3D 标量场计算给定等值面的 2D 表面网格。
定义输出网格的三角形和顶点被写入
indices和verts数组。- 参数：
- Raises:
ValueError –
field不是 3D 数组。ValueError – 行进立方的形状与
field的形状不匹配。RuntimeError –
max_verts和/或max_tris可能太小，无法容纳表面网格。
- 返回类型：
None
性能分析#
wp.ScopedTimer 对象可用于对 Warp 应用程序的性能获得一些基本洞察
with wp.ScopedTimer("grid build"):
self.grid.build(self.x, self.point_radius)
这会在运行时向标准输出流打印如下内容
grid build took 0.06 ms
有关更多信息，请参见性能分析文档。
- class warp.ScopedTimer(
- name,
- active=True,
- print=True,
- detailed=False,
- dict=None,
- use_nvtx=False,
- color='rapids',
- synchronize=False,
- cuda_filter=0,
- report_func=None,
- skip_tape=False,
- 参数：
- indent = -1
- enabled = True
- __init__(
- name,
- active=True,
- print=True,
- detailed=False,
- dict=None,
- use_nvtx=False,
- color='rapids',
- synchronize=False,
- cuda_filter=0,
- report_func=None,
- skip_tape=False,
计时器的上下文管理器对象
- 参数：
name (str) – 计时器名称
active (bool) – 启用此计时器
print (bool) – 在上下文管理器退出时，将经过时间打印到
sys.stdoutdetailed (bool) – 使用 cProfile 收集额外的性能分析数据，并在上下文退出时调用
print_stats()dict (Dict[str, List[float]] | None) – 一个字典列表，将使用
name作为键将经过时间附加到其中use_nvtx (bool) – 如果为 true，计时功能将被 NVTX 范围取代
color (int | str) – 与 NVTX 范围关联的 ARGB 值（例如 0x00FFFF）或颜色名称（例如 'cyan'）
synchronize (bool) – 使 CPU 线程与任何未完成的 CUDA 工作同步，以返回准确的 GPU 计时
cuda_filter (int) – 用于 CUDA 活动计时的过滤标志，例如
warp.TIMING_KERNEL或warp.TIMING_ALLreport_func (Callable[[List[TimingResult], str], None] | None) – 用于打印活动报告的回调函数。如果为
None，将使用wp.timing_print()。skip_tape (bool) – 如果为 true，计时器将不会记录在 tape 中
- extra_msg
可设置为在上下文退出时添加到打印输出中的字符串。
- 类型：
- elapsed
与此对象一起使用的
with块的持续时间- 类型：
- timing_results
活动计时结果列表，如果使用
cuda_filter请求收集- 类型：
List[TimingResult]
进程间通信 (IPC)#
进程间通信可用于在进程之间共享 Warp 数组和事件，而无需创建底层数据的副本。
使用 IPC 的一些基本要求包括
Linux 操作系统（但请注意，NVIDIA Jetson 等集成设备不支持 CUDA IPC）
数组必须使用默认内存分配器分配在 GPU 设备上（参见分配器）
wp.ScopedMempool上下文管理器对于暂时禁用内存池以分配可使用 IPC 共享的数组非常有用。
设备上对 IPC 的支持由 Device 的 is_ipc_supported 属性指示。如果 Warp 库已使用 CUDA 11 编译，此设备属性将为 None，表示无法使用 CUDA API 确定 IPC 支持。
要在进程之间共享 Warp 数组，请在发起进程中使用 array.ipc_handle() 获取数组内存分配的 IPC 句柄。该句柄是一个长度为 64 的 bytes 对象。IPC 句柄以及数组信息（数据类型、形状，以及可选的步长）应与另一个进程共享，例如通过共享内存或文件。另一个进程可以使用此信息通过调用 from_ipc_handle() 来导入原始数组。
事件也可以以类似的方式共享，但它们必须使用 interprocess=True 构建。此外，不能同时创建具有 interprocess=True 和 enable_timing=True 的事件。在发起进程中使用 Event.ipc_handle() 获取事件的 IPC 句柄。另一个进程可以使用此信息通过调用 event_from_ipc_handle() 来导入原始事件。
- warp.from_ipc_handle(handle, dtype, shape, strides=None, device=None)[source]#
从 IPC 句柄创建数组。
dtype、shape和可选的strides参数应与创建handle的array中的值匹配。- 参数：
- 返回值：
从由进程间内存句柄
handle描述的现有内存分配创建的数组。不复制底层数据。对数组数据的修改将反映在导出
handle的原始进程中。- Raises:
RuntimeError –
device不支持 IPC。- 返回类型：
- warp.event_from_ipc_handle(handle, device=None)[source]#
从 IPC 句柄创建事件。
- 参数：
handle – 现有 CUDA 事件的进程间事件句柄。
device (Devicelike) – 与数组关联的设备。
- 返回值：
从进程间事件句柄
handle创建的事件。- Raises:
RuntimeError –
device不支持 IPC。- 返回类型：
LTO 缓存#
MathDx 为 GEMM、Cholesky 和 FFT 块操作生成链接时优化 (LTO) 文件。Warp 缓存这些文件以加快核函数编译速度。每个 LTO 文件对应于特定的线性代数求解器配置，并且与调用其相应例程的核函数无关。因此，LTO 存储在独立于给定模块核函数缓存的缓存中，即使调用 wp.clear_kernel_cache()，LTO 也将保持缓存状态。wp.clear_lto_cache() 可用于清除 LTO 缓存。