warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__py3-none-manylinux_2_34_aarch64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of warp-lang might be problematic. Click here for more details.
- warp/__init__.py +3 -1
- warp/__init__.pyi +3489 -1
- warp/autograd.py +45 -122
- warp/bin/warp-clang.so +0 -0
- warp/bin/warp.so +0 -0
- warp/build.py +241 -252
- warp/build_dll.py +125 -26
- warp/builtins.py +1907 -384
- warp/codegen.py +257 -101
- warp/config.py +12 -1
- warp/constants.py +1 -1
- warp/context.py +657 -223
- warp/dlpack.py +1 -1
- warp/examples/benchmarks/benchmark_cloth.py +2 -2
- warp/examples/benchmarks/benchmark_tile_sort.py +155 -0
- warp/examples/core/example_sample_mesh.py +1 -1
- warp/examples/core/example_spin_lock.py +93 -0
- warp/examples/core/example_work_queue.py +118 -0
- warp/examples/fem/example_adaptive_grid.py +5 -5
- warp/examples/fem/example_apic_fluid.py +1 -1
- warp/examples/fem/example_burgers.py +1 -1
- warp/examples/fem/example_convection_diffusion.py +9 -6
- warp/examples/fem/example_darcy_ls_optimization.py +489 -0
- warp/examples/fem/example_deformed_geometry.py +1 -1
- warp/examples/fem/example_diffusion.py +2 -2
- warp/examples/fem/example_diffusion_3d.py +1 -1
- warp/examples/fem/example_distortion_energy.py +1 -1
- warp/examples/fem/example_elastic_shape_optimization.py +387 -0
- warp/examples/fem/example_magnetostatics.py +5 -3
- warp/examples/fem/example_mixed_elasticity.py +5 -3
- warp/examples/fem/example_navier_stokes.py +11 -9
- warp/examples/fem/example_nonconforming_contact.py +5 -3
- warp/examples/fem/example_streamlines.py +8 -3
- warp/examples/fem/utils.py +9 -8
- warp/examples/interop/example_jax_ffi_callback.py +2 -2
- warp/examples/optim/example_drone.py +1 -1
- warp/examples/sim/example_cloth.py +1 -1
- warp/examples/sim/example_cloth_self_contact.py +48 -54
- warp/examples/tile/example_tile_block_cholesky.py +502 -0
- warp/examples/tile/example_tile_cholesky.py +2 -1
- warp/examples/tile/example_tile_convolution.py +1 -1
- warp/examples/tile/example_tile_filtering.py +1 -1
- warp/examples/tile/example_tile_matmul.py +1 -1
- warp/examples/tile/example_tile_mlp.py +2 -0
- warp/fabric.py +7 -7
- warp/fem/__init__.py +5 -0
- warp/fem/adaptivity.py +1 -1
- warp/fem/cache.py +152 -63
- warp/fem/dirichlet.py +2 -2
- warp/fem/domain.py +136 -6
- warp/fem/field/field.py +141 -99
- warp/fem/field/nodal_field.py +85 -39
- warp/fem/field/virtual.py +97 -52
- warp/fem/geometry/adaptive_nanogrid.py +91 -86
- warp/fem/geometry/closest_point.py +13 -0
- warp/fem/geometry/deformed_geometry.py +102 -40
- warp/fem/geometry/element.py +56 -2
- warp/fem/geometry/geometry.py +323 -22
- warp/fem/geometry/grid_2d.py +157 -62
- warp/fem/geometry/grid_3d.py +116 -20
- warp/fem/geometry/hexmesh.py +86 -20
- warp/fem/geometry/nanogrid.py +166 -86
- warp/fem/geometry/partition.py +59 -25
- warp/fem/geometry/quadmesh.py +86 -135
- warp/fem/geometry/tetmesh.py +47 -119
- warp/fem/geometry/trimesh.py +77 -270
- warp/fem/integrate.py +107 -52
- warp/fem/linalg.py +25 -58
- warp/fem/operator.py +124 -27
- warp/fem/quadrature/pic_quadrature.py +36 -14
- warp/fem/quadrature/quadrature.py +40 -16
- warp/fem/space/__init__.py +1 -1
- warp/fem/space/basis_function_space.py +66 -46
- warp/fem/space/basis_space.py +17 -4
- warp/fem/space/dof_mapper.py +1 -1
- warp/fem/space/function_space.py +2 -2
- warp/fem/space/grid_2d_function_space.py +4 -1
- warp/fem/space/hexmesh_function_space.py +4 -2
- warp/fem/space/nanogrid_function_space.py +3 -1
- warp/fem/space/partition.py +11 -2
- warp/fem/space/quadmesh_function_space.py +4 -1
- warp/fem/space/restriction.py +5 -2
- warp/fem/space/shape/__init__.py +10 -8
- warp/fem/space/tetmesh_function_space.py +4 -1
- warp/fem/space/topology.py +52 -21
- warp/fem/space/trimesh_function_space.py +4 -1
- warp/fem/utils.py +53 -8
- warp/jax.py +1 -2
- warp/jax_experimental/ffi.py +12 -17
- warp/jax_experimental/xla_ffi.py +37 -24
- warp/math.py +171 -1
- warp/native/array.h +99 -0
- warp/native/builtin.h +174 -31
- warp/native/coloring.cpp +1 -1
- warp/native/exports.h +118 -63
- warp/native/intersect.h +3 -3
- warp/native/mat.h +5 -10
- warp/native/mathdx.cpp +11 -5
- warp/native/matnn.h +1 -123
- warp/native/quat.h +28 -4
- warp/native/sparse.cpp +121 -258
- warp/native/sparse.cu +181 -274
- warp/native/spatial.h +305 -17
- warp/native/tile.h +583 -72
- warp/native/tile_radix_sort.h +1108 -0
- warp/native/tile_reduce.h +237 -2
- warp/native/tile_scan.h +240 -0
- warp/native/tuple.h +189 -0
- warp/native/vec.h +6 -16
- warp/native/warp.cpp +36 -4
- warp/native/warp.cu +574 -51
- warp/native/warp.h +47 -74
- warp/optim/linear.py +5 -1
- warp/paddle.py +7 -8
- warp/py.typed +0 -0
- warp/render/render_opengl.py +58 -29
- warp/render/render_usd.py +124 -61
- warp/sim/__init__.py +9 -0
- warp/sim/collide.py +252 -78
- warp/sim/graph_coloring.py +8 -1
- warp/sim/import_mjcf.py +4 -3
- warp/sim/import_usd.py +11 -7
- warp/sim/integrator.py +5 -2
- warp/sim/integrator_euler.py +1 -1
- warp/sim/integrator_featherstone.py +1 -1
- warp/sim/integrator_vbd.py +751 -320
- warp/sim/integrator_xpbd.py +1 -1
- warp/sim/model.py +265 -260
- warp/sim/utils.py +10 -7
- warp/sparse.py +303 -166
- warp/tape.py +52 -51
- warp/tests/cuda/test_conditional_captures.py +1046 -0
- warp/tests/cuda/test_streams.py +1 -1
- warp/tests/geometry/test_volume.py +2 -2
- warp/tests/interop/test_dlpack.py +9 -9
- warp/tests/interop/test_jax.py +0 -1
- warp/tests/run_coverage_serial.py +1 -1
- warp/tests/sim/disabled_kinematics.py +2 -2
- warp/tests/sim/{test_vbd.py → test_cloth.py} +296 -113
- warp/tests/sim/test_collision.py +159 -51
- warp/tests/sim/test_coloring.py +15 -1
- warp/tests/test_array.py +254 -2
- warp/tests/test_array_reduce.py +2 -2
- warp/tests/test_atomic_cas.py +299 -0
- warp/tests/test_codegen.py +142 -19
- warp/tests/test_conditional.py +47 -1
- warp/tests/test_ctypes.py +0 -20
- warp/tests/test_devices.py +8 -0
- warp/tests/test_fabricarray.py +4 -2
- warp/tests/test_fem.py +58 -25
- warp/tests/test_func.py +42 -1
- warp/tests/test_grad.py +1 -1
- warp/tests/test_lerp.py +1 -3
- warp/tests/test_map.py +481 -0
- warp/tests/test_mat.py +1 -24
- warp/tests/test_quat.py +6 -15
- warp/tests/test_rounding.py +10 -38
- warp/tests/test_runlength_encode.py +7 -7
- warp/tests/test_smoothstep.py +1 -1
- warp/tests/test_sparse.py +51 -2
- warp/tests/test_spatial.py +507 -1
- warp/tests/test_struct.py +2 -2
- warp/tests/test_tuple.py +265 -0
- warp/tests/test_types.py +2 -2
- warp/tests/test_utils.py +24 -18
- warp/tests/tile/test_tile.py +420 -1
- warp/tests/tile/test_tile_mathdx.py +518 -14
- warp/tests/tile/test_tile_reduce.py +213 -0
- warp/tests/tile/test_tile_shared_memory.py +130 -1
- warp/tests/tile/test_tile_sort.py +117 -0
- warp/tests/unittest_suites.py +4 -6
- warp/types.py +462 -308
- warp/utils.py +647 -86
- {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/METADATA +20 -6
- {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/RECORD +178 -166
- warp/stubs.py +0 -3381
- warp/tests/sim/test_xpbd.py +0 -399
- warp/tests/test_mlp.py +0 -282
- {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/WHEEL +0 -0
- {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/licenses/LICENSE.md +0 -0
- {warp_lang-1.7.2rc1.dist-info → warp_lang-1.8.0.dist-info}/top_level.txt +0 -0
warp/__init__.pyi
CHANGED
|
@@ -1 +1,3489 @@
|
|
|
1
|
-
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
3
|
+
#
|
|
4
|
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
5
|
+
# you may not use this file except in compliance with the License.
|
|
6
|
+
# You may obtain a copy of the License at
|
|
7
|
+
#
|
|
8
|
+
# http://www.apache.org/licenses/LICENSE-2.0
|
|
9
|
+
#
|
|
10
|
+
# Unless required by applicable law or agreed to in writing, software
|
|
11
|
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
12
|
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
13
|
+
# See the License for the specific language governing permissions and
|
|
14
|
+
# limitations under the License.
|
|
15
|
+
|
|
16
|
+
# Autogenerated file, do not edit, this file provides stubs for builtins autocomplete in VSCode, PyCharm, etc
|
|
17
|
+
|
|
18
|
+
from typing import Any
|
|
19
|
+
from typing import Tuple
|
|
20
|
+
from typing import Callable
|
|
21
|
+
from typing import TypeVar
|
|
22
|
+
from typing import Generic
|
|
23
|
+
from typing import overload as over
|
|
24
|
+
|
|
25
|
+
Length = TypeVar("Length", bound=int)
|
|
26
|
+
Rows = TypeVar("Rows", bound=int)
|
|
27
|
+
Cols = TypeVar("Cols", bound=int)
|
|
28
|
+
DType = TypeVar("DType")
|
|
29
|
+
Shape = TypeVar("Shape")
|
|
30
|
+
Vector = Generic[Length, Scalar]
|
|
31
|
+
Matrix = Generic[Rows, Cols, Scalar]
|
|
32
|
+
Quaternion = Generic[Float]
|
|
33
|
+
Transformation = Generic[Float]
|
|
34
|
+
Array = Generic[DType]
|
|
35
|
+
FabricArray = Generic[DType]
|
|
36
|
+
IndexedFabricArray = Generic[DType]
|
|
37
|
+
Tile = Generic[DType, Shape]
|
|
38
|
+
|
|
39
|
+
from warp.types import array, array1d, array2d, array3d, array4d, constant, from_ptr
|
|
40
|
+
from warp.types import indexedarray, indexedarray1d, indexedarray2d, indexedarray3d, indexedarray4d
|
|
41
|
+
from warp.fabric import fabricarray, fabricarrayarray, indexedfabricarray, indexedfabricarrayarray
|
|
42
|
+
from warp.types import tile
|
|
43
|
+
|
|
44
|
+
from warp.types import bool, int8, uint8, int16, uint16, int32, uint32, int64, uint64, float16, float32, float64
|
|
45
|
+
from warp.types import vec2, vec2b, vec2ub, vec2s, vec2us, vec2i, vec2ui, vec2l, vec2ul, vec2h, vec2f, vec2d
|
|
46
|
+
from warp.types import vec3, vec3b, vec3ub, vec3s, vec3us, vec3i, vec3ui, vec3l, vec3ul, vec3h, vec3f, vec3d
|
|
47
|
+
from warp.types import vec4, vec4b, vec4ub, vec4s, vec4us, vec4i, vec4ui, vec4l, vec4ul, vec4h, vec4f, vec4d
|
|
48
|
+
from warp.types import mat22, mat22h, mat22f, mat22d
|
|
49
|
+
from warp.types import mat33, mat33h, mat33f, mat33d
|
|
50
|
+
from warp.types import mat44, mat44h, mat44f, mat44d
|
|
51
|
+
from warp.types import quat, quath, quatf, quatd
|
|
52
|
+
from warp.types import transform, transformh, transformf, transformd
|
|
53
|
+
from warp.types import spatial_vector, spatial_vectorh, spatial_vectorf, spatial_vectord
|
|
54
|
+
from warp.types import spatial_matrix, spatial_matrixh, spatial_matrixf, spatial_matrixd
|
|
55
|
+
|
|
56
|
+
from warp.types import Int, Float, Scalar
|
|
57
|
+
|
|
58
|
+
from warp.types import Bvh, Mesh, HashGrid, Volume, MarchingCubes
|
|
59
|
+
from warp.types import BvhQuery, HashGridQuery, MeshQueryAABB, MeshQueryPoint, MeshQueryRay
|
|
60
|
+
|
|
61
|
+
from warp.types import matmul, adj_matmul, batched_matmul, adj_batched_matmul
|
|
62
|
+
|
|
63
|
+
from warp.types import vector as vec
|
|
64
|
+
from warp.types import matrix as mat
|
|
65
|
+
|
|
66
|
+
from warp.types import dtype_from_numpy, dtype_to_numpy
|
|
67
|
+
|
|
68
|
+
from warp.types import from_ipc_handle
|
|
69
|
+
|
|
70
|
+
from warp.context import init, func, func_grad, func_replay, func_native, kernel, struct, overload
|
|
71
|
+
from warp.context import is_cpu_available, is_cuda_available, is_device_available
|
|
72
|
+
from warp.context import get_devices, get_preferred_device
|
|
73
|
+
from warp.context import get_cuda_devices, get_cuda_device_count, get_cuda_device, map_cuda_device, unmap_cuda_device
|
|
74
|
+
from warp.context import get_device, set_device, synchronize_device
|
|
75
|
+
from warp.context import (
|
|
76
|
+
zeros,
|
|
77
|
+
zeros_like,
|
|
78
|
+
ones,
|
|
79
|
+
ones_like,
|
|
80
|
+
full,
|
|
81
|
+
full_like,
|
|
82
|
+
clone,
|
|
83
|
+
empty,
|
|
84
|
+
empty_like,
|
|
85
|
+
copy,
|
|
86
|
+
from_numpy,
|
|
87
|
+
launch,
|
|
88
|
+
launch_tiled,
|
|
89
|
+
synchronize,
|
|
90
|
+
force_load,
|
|
91
|
+
load_module,
|
|
92
|
+
event_from_ipc_handle,
|
|
93
|
+
)
|
|
94
|
+
from warp.context import set_module_options, get_module_options, get_module
|
|
95
|
+
from warp.context import capture_begin, capture_end, capture_launch, capture_if, capture_while, capture_debug_dot_print
|
|
96
|
+
from warp.context import Kernel, Function, Launch
|
|
97
|
+
from warp.context import Stream, get_stream, set_stream, wait_stream, synchronize_stream
|
|
98
|
+
from warp.context import Event, record_event, wait_event, synchronize_event, get_event_elapsed_time
|
|
99
|
+
from warp.context import RegisteredGLBuffer
|
|
100
|
+
from warp.context import is_mempool_supported, is_mempool_enabled, set_mempool_enabled
|
|
101
|
+
from warp.context import (
|
|
102
|
+
set_mempool_release_threshold,
|
|
103
|
+
get_mempool_release_threshold,
|
|
104
|
+
get_mempool_used_mem_current,
|
|
105
|
+
get_mempool_used_mem_high,
|
|
106
|
+
)
|
|
107
|
+
from warp.context import is_mempool_access_supported, is_mempool_access_enabled, set_mempool_access_enabled
|
|
108
|
+
from warp.context import is_peer_access_supported, is_peer_access_enabled, set_peer_access_enabled
|
|
109
|
+
|
|
110
|
+
from warp.tape import Tape
|
|
111
|
+
from warp.utils import ScopedTimer, ScopedDevice, ScopedStream
|
|
112
|
+
from warp.utils import ScopedMempool, ScopedMempoolAccess, ScopedPeerAccess
|
|
113
|
+
from warp.utils import ScopedCapture
|
|
114
|
+
from warp.utils import transform_expand, quat_between_vectors
|
|
115
|
+
from warp.utils import TimingResult, timing_begin, timing_end, timing_print
|
|
116
|
+
from warp.utils import (
|
|
117
|
+
TIMING_KERNEL,
|
|
118
|
+
TIMING_KERNEL_BUILTIN,
|
|
119
|
+
TIMING_MEMCPY,
|
|
120
|
+
TIMING_MEMSET,
|
|
121
|
+
TIMING_GRAPH,
|
|
122
|
+
TIMING_ALL,
|
|
123
|
+
)
|
|
124
|
+
from warp.utils import map
|
|
125
|
+
|
|
126
|
+
from warp.torch import from_torch, to_torch
|
|
127
|
+
from warp.torch import dtype_from_torch, dtype_to_torch
|
|
128
|
+
from warp.torch import device_from_torch, device_to_torch
|
|
129
|
+
from warp.torch import stream_from_torch, stream_to_torch
|
|
130
|
+
|
|
131
|
+
from warp.jax import from_jax, to_jax
|
|
132
|
+
from warp.jax import dtype_from_jax, dtype_to_jax
|
|
133
|
+
from warp.jax import device_from_jax, device_to_jax
|
|
134
|
+
|
|
135
|
+
from warp.dlpack import from_dlpack, to_dlpack
|
|
136
|
+
|
|
137
|
+
from warp.paddle import from_paddle, to_paddle
|
|
138
|
+
from warp.paddle import dtype_from_paddle, dtype_to_paddle
|
|
139
|
+
from warp.paddle import device_from_paddle, device_to_paddle
|
|
140
|
+
from warp.paddle import stream_from_paddle
|
|
141
|
+
|
|
142
|
+
from warp.build import clear_kernel_cache
|
|
143
|
+
from warp.build import clear_lto_cache
|
|
144
|
+
|
|
145
|
+
from warp.constants import *
|
|
146
|
+
|
|
147
|
+
from . import builtins
|
|
148
|
+
from warp.builtins import static
|
|
149
|
+
|
|
150
|
+
from warp.math import *
|
|
151
|
+
|
|
152
|
+
import warp.config as config
|
|
153
|
+
|
|
154
|
+
__version__ = config.version
|
|
155
|
+
|
|
156
|
+
@over
|
|
157
|
+
def min(a: Scalar, b: Scalar) -> Scalar:
|
|
158
|
+
"""Return the minimum of two scalars."""
|
|
159
|
+
...
|
|
160
|
+
|
|
161
|
+
@over
|
|
162
|
+
def min(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
163
|
+
"""Return the element-wise minimum of two vectors."""
|
|
164
|
+
...
|
|
165
|
+
|
|
166
|
+
@over
|
|
167
|
+
def min(a: Vector[Any, Scalar]) -> Scalar:
|
|
168
|
+
"""Return the minimum element of a vector ``a``."""
|
|
169
|
+
...
|
|
170
|
+
|
|
171
|
+
@over
|
|
172
|
+
def max(a: Scalar, b: Scalar) -> Scalar:
|
|
173
|
+
"""Return the maximum of two scalars."""
|
|
174
|
+
...
|
|
175
|
+
|
|
176
|
+
@over
|
|
177
|
+
def max(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
178
|
+
"""Return the element-wise maximum of two vectors."""
|
|
179
|
+
...
|
|
180
|
+
|
|
181
|
+
@over
|
|
182
|
+
def max(a: Vector[Any, Scalar]) -> Scalar:
|
|
183
|
+
"""Return the maximum element of a vector ``a``."""
|
|
184
|
+
...
|
|
185
|
+
|
|
186
|
+
@over
|
|
187
|
+
def clamp(x: Scalar, low: Scalar, high: Scalar) -> Scalar:
|
|
188
|
+
"""Clamp the value of ``x`` to the range [low, high]."""
|
|
189
|
+
...
|
|
190
|
+
|
|
191
|
+
@over
|
|
192
|
+
def abs(x: Scalar) -> Scalar:
|
|
193
|
+
"""Return the absolute value of ``x``."""
|
|
194
|
+
...
|
|
195
|
+
|
|
196
|
+
@over
|
|
197
|
+
def abs(x: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
198
|
+
"""Return the absolute values of the elements of ``x``."""
|
|
199
|
+
...
|
|
200
|
+
|
|
201
|
+
@over
|
|
202
|
+
def sign(x: Scalar) -> Scalar:
|
|
203
|
+
"""Return -1 if ``x`` < 0, return 1 otherwise."""
|
|
204
|
+
...
|
|
205
|
+
|
|
206
|
+
@over
|
|
207
|
+
def sign(x: Vector[Any, Scalar]) -> Scalar:
|
|
208
|
+
"""Return -1 for the negative elements of ``x``, and 1 otherwise."""
|
|
209
|
+
...
|
|
210
|
+
|
|
211
|
+
@over
|
|
212
|
+
def step(x: Scalar) -> Scalar:
|
|
213
|
+
"""Return 1.0 if ``x`` < 0.0, return 0.0 otherwise."""
|
|
214
|
+
...
|
|
215
|
+
|
|
216
|
+
@over
|
|
217
|
+
def nonzero(x: Scalar) -> Scalar:
|
|
218
|
+
"""Return 1.0 if ``x`` is not equal to zero, return 0.0 otherwise."""
|
|
219
|
+
...
|
|
220
|
+
|
|
221
|
+
@over
|
|
222
|
+
def sin(x: Float) -> Float:
|
|
223
|
+
"""Return the sine of ``x`` in radians."""
|
|
224
|
+
...
|
|
225
|
+
|
|
226
|
+
@over
|
|
227
|
+
def cos(x: Float) -> Float:
|
|
228
|
+
"""Return the cosine of ``x`` in radians."""
|
|
229
|
+
...
|
|
230
|
+
|
|
231
|
+
@over
|
|
232
|
+
def acos(x: Float) -> Float:
|
|
233
|
+
"""Return arccos of ``x`` in radians. Inputs are automatically clamped to [-1.0, 1.0]."""
|
|
234
|
+
...
|
|
235
|
+
|
|
236
|
+
@over
|
|
237
|
+
def asin(x: Float) -> Float:
|
|
238
|
+
"""Return arcsin of ``x`` in radians. Inputs are automatically clamped to [-1.0, 1.0]."""
|
|
239
|
+
...
|
|
240
|
+
|
|
241
|
+
@over
|
|
242
|
+
def sqrt(x: Float) -> Float:
|
|
243
|
+
"""Return the square root of ``x``, where ``x`` is positive."""
|
|
244
|
+
...
|
|
245
|
+
|
|
246
|
+
@over
|
|
247
|
+
def cbrt(x: Float) -> Float:
|
|
248
|
+
"""Return the cube root of ``x``."""
|
|
249
|
+
...
|
|
250
|
+
|
|
251
|
+
@over
|
|
252
|
+
def tan(x: Float) -> Float:
|
|
253
|
+
"""Return the tangent of ``x`` in radians."""
|
|
254
|
+
...
|
|
255
|
+
|
|
256
|
+
@over
|
|
257
|
+
def atan(x: Float) -> Float:
|
|
258
|
+
"""Return the arctangent of ``x`` in radians."""
|
|
259
|
+
...
|
|
260
|
+
|
|
261
|
+
@over
|
|
262
|
+
def atan2(y: Float, x: Float) -> Float:
|
|
263
|
+
"""Return the 2-argument arctangent, atan2, of the point ``(x, y)`` in radians."""
|
|
264
|
+
...
|
|
265
|
+
|
|
266
|
+
@over
|
|
267
|
+
def sinh(x: Float) -> Float:
|
|
268
|
+
"""Return the sinh of ``x``."""
|
|
269
|
+
...
|
|
270
|
+
|
|
271
|
+
@over
|
|
272
|
+
def cosh(x: Float) -> Float:
|
|
273
|
+
"""Return the cosh of ``x``."""
|
|
274
|
+
...
|
|
275
|
+
|
|
276
|
+
@over
|
|
277
|
+
def tanh(x: Float) -> Float:
|
|
278
|
+
"""Return the tanh of ``x``."""
|
|
279
|
+
...
|
|
280
|
+
|
|
281
|
+
@over
|
|
282
|
+
def degrees(x: Float) -> Float:
|
|
283
|
+
"""Convert ``x`` from radians into degrees."""
|
|
284
|
+
...
|
|
285
|
+
|
|
286
|
+
@over
|
|
287
|
+
def radians(x: Float) -> Float:
|
|
288
|
+
"""Convert ``x`` from degrees into radians."""
|
|
289
|
+
...
|
|
290
|
+
|
|
291
|
+
@over
|
|
292
|
+
def log(x: Float) -> Float:
|
|
293
|
+
"""Return the natural logarithm (base-e) of ``x``, where ``x`` is positive."""
|
|
294
|
+
...
|
|
295
|
+
|
|
296
|
+
@over
|
|
297
|
+
def log2(x: Float) -> Float:
|
|
298
|
+
"""Return the binary logarithm (base-2) of ``x``, where ``x`` is positive."""
|
|
299
|
+
...
|
|
300
|
+
|
|
301
|
+
@over
|
|
302
|
+
def log10(x: Float) -> Float:
|
|
303
|
+
"""Return the common logarithm (base-10) of ``x``, where ``x`` is positive."""
|
|
304
|
+
...
|
|
305
|
+
|
|
306
|
+
@over
|
|
307
|
+
def exp(x: Float) -> Float:
|
|
308
|
+
"""Return the value of the exponential function :math:`e^x`."""
|
|
309
|
+
...
|
|
310
|
+
|
|
311
|
+
@over
|
|
312
|
+
def pow(x: Float, y: Float) -> Float:
|
|
313
|
+
"""Return the result of ``x`` raised to power of ``y``."""
|
|
314
|
+
...
|
|
315
|
+
|
|
316
|
+
@over
|
|
317
|
+
def round(x: Float) -> Float:
|
|
318
|
+
"""Return the nearest integer value to ``x``, rounding halfway cases away from zero.
|
|
319
|
+
|
|
320
|
+
This is the most intuitive form of rounding in the colloquial sense, but can be slower than other options like :func:`warp.rint()`.
|
|
321
|
+
Differs from :func:`numpy.round()`, which behaves the same way as :func:`numpy.rint()`.
|
|
322
|
+
"""
|
|
323
|
+
...
|
|
324
|
+
|
|
325
|
+
@over
|
|
326
|
+
def rint(x: Float) -> Float:
|
|
327
|
+
"""Return the nearest integer value to ``x``, rounding halfway cases to nearest even integer.
|
|
328
|
+
|
|
329
|
+
It is generally faster than :func:`warp.round()`. Equivalent to :func:`numpy.rint()`.
|
|
330
|
+
"""
|
|
331
|
+
...
|
|
332
|
+
|
|
333
|
+
@over
|
|
334
|
+
def trunc(x: Float) -> Float:
|
|
335
|
+
"""Return the nearest integer that is closer to zero than ``x``.
|
|
336
|
+
|
|
337
|
+
In other words, it discards the fractional part of ``x``.
|
|
338
|
+
It is similar to casting ``float(int(a))``, but preserves the negative sign when ``x`` is in the range [-0.0, -1.0).
|
|
339
|
+
Equivalent to :func:`numpy.trunc()` and :func:`numpy.fix()`.
|
|
340
|
+
"""
|
|
341
|
+
...
|
|
342
|
+
|
|
343
|
+
@over
|
|
344
|
+
def floor(x: Float) -> Float:
|
|
345
|
+
"""Return the largest integer that is less than or equal to ``x``."""
|
|
346
|
+
...
|
|
347
|
+
|
|
348
|
+
@over
|
|
349
|
+
def ceil(x: Float) -> Float:
|
|
350
|
+
"""Return the smallest integer that is greater than or equal to ``x``."""
|
|
351
|
+
...
|
|
352
|
+
|
|
353
|
+
@over
|
|
354
|
+
def frac(x: Float) -> Float:
|
|
355
|
+
"""Retrieve the fractional part of ``x``.
|
|
356
|
+
|
|
357
|
+
In other words, it discards the integer part of ``x`` and is equivalent to ``x - trunc(x)``.
|
|
358
|
+
"""
|
|
359
|
+
...
|
|
360
|
+
|
|
361
|
+
@over
|
|
362
|
+
def isfinite(a: Scalar) -> bool:
|
|
363
|
+
"""Return ``True`` if ``a`` is a finite number, otherwise return ``False``."""
|
|
364
|
+
...
|
|
365
|
+
|
|
366
|
+
@over
|
|
367
|
+
def isfinite(a: Vector[Any, Scalar]) -> bool:
|
|
368
|
+
"""Return ``True`` if all elements of the vector ``a`` are finite, otherwise return ``False``."""
|
|
369
|
+
...
|
|
370
|
+
|
|
371
|
+
@over
|
|
372
|
+
def isfinite(a: Quaternion[Scalar]) -> bool:
|
|
373
|
+
"""Return ``True`` if all elements of the quaternion ``a`` are finite, otherwise return ``False``."""
|
|
374
|
+
...
|
|
375
|
+
|
|
376
|
+
@over
|
|
377
|
+
def isfinite(a: Matrix[Any, Any, Scalar]) -> bool:
|
|
378
|
+
"""Return ``True`` if all elements of the matrix ``a`` are finite, otherwise return ``False``."""
|
|
379
|
+
...
|
|
380
|
+
|
|
381
|
+
@over
|
|
382
|
+
def isnan(a: Scalar) -> bool:
|
|
383
|
+
"""Return ``True`` if ``a`` is NaN, otherwise return ``False``."""
|
|
384
|
+
...
|
|
385
|
+
|
|
386
|
+
@over
|
|
387
|
+
def isnan(a: Vector[Any, Scalar]) -> bool:
|
|
388
|
+
"""Return ``True`` if any element of the vector ``a`` is NaN, otherwise return ``False``."""
|
|
389
|
+
...
|
|
390
|
+
|
|
391
|
+
@over
|
|
392
|
+
def isnan(a: Quaternion[Scalar]) -> bool:
|
|
393
|
+
"""Return ``True`` if any element of the quaternion ``a`` is NaN, otherwise return ``False``."""
|
|
394
|
+
...
|
|
395
|
+
|
|
396
|
+
@over
|
|
397
|
+
def isnan(a: Matrix[Any, Any, Scalar]) -> bool:
|
|
398
|
+
"""Return ``True`` if any element of the matrix ``a`` is NaN, otherwise return ``False``."""
|
|
399
|
+
...
|
|
400
|
+
|
|
401
|
+
@over
|
|
402
|
+
def isinf(a: Scalar) -> bool:
|
|
403
|
+
"""Return ``True`` if ``a`` is positive or negative infinity, otherwise return ``False``."""
|
|
404
|
+
...
|
|
405
|
+
|
|
406
|
+
@over
|
|
407
|
+
def isinf(a: Vector[Any, Scalar]) -> bool:
|
|
408
|
+
"""Return ``True`` if any element of the vector ``a`` is positive or negative infinity, otherwise return ``False``."""
|
|
409
|
+
...
|
|
410
|
+
|
|
411
|
+
@over
|
|
412
|
+
def isinf(a: Quaternion[Scalar]) -> bool:
|
|
413
|
+
"""Return ``True`` if any element of the quaternion ``a`` is positive or negative infinity, otherwise return ``False``."""
|
|
414
|
+
...
|
|
415
|
+
|
|
416
|
+
@over
|
|
417
|
+
def isinf(a: Matrix[Any, Any, Scalar]) -> bool:
|
|
418
|
+
"""Return ``True`` if any element of the matrix ``a`` is positive or negative infinity, otherwise return ``False``."""
|
|
419
|
+
...
|
|
420
|
+
|
|
421
|
+
@over
|
|
422
|
+
def dot(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Scalar:
|
|
423
|
+
"""Compute the dot product between two vectors."""
|
|
424
|
+
...
|
|
425
|
+
|
|
426
|
+
@over
|
|
427
|
+
def dot(a: Quaternion[Float], b: Quaternion[Float]) -> Float:
|
|
428
|
+
"""Compute the dot product between two quaternions."""
|
|
429
|
+
...
|
|
430
|
+
|
|
431
|
+
@over
|
|
432
|
+
def ddot(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Scalar:
|
|
433
|
+
"""Compute the double dot product between two matrices."""
|
|
434
|
+
...
|
|
435
|
+
|
|
436
|
+
@over
|
|
437
|
+
def argmin(a: Vector[Any, Scalar]) -> uint32:
|
|
438
|
+
"""Return the index of the minimum element of a vector ``a``."""
|
|
439
|
+
...
|
|
440
|
+
|
|
441
|
+
@over
|
|
442
|
+
def argmax(a: Vector[Any, Scalar]) -> uint32:
|
|
443
|
+
"""Return the index of the maximum element of a vector ``a``."""
|
|
444
|
+
...
|
|
445
|
+
|
|
446
|
+
@over
|
|
447
|
+
def outer(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
448
|
+
"""Compute the outer product ``a*b^T`` for two vectors."""
|
|
449
|
+
...
|
|
450
|
+
|
|
451
|
+
@over
|
|
452
|
+
def cross(a: Vector[3, Scalar], b: Vector[3, Scalar]) -> Vector[3, Scalar]:
|
|
453
|
+
"""Compute the cross product of two 3D vectors."""
|
|
454
|
+
...
|
|
455
|
+
|
|
456
|
+
@over
|
|
457
|
+
def skew(vec: Vector[3, Scalar]) -> Matrix[3, 3, Scalar]:
|
|
458
|
+
"""Compute the skew-symmetric 3x3 matrix for a 3D vector ``vec``."""
|
|
459
|
+
...
|
|
460
|
+
|
|
461
|
+
@over
|
|
462
|
+
def length(a: Vector[Any, Float]) -> Float:
|
|
463
|
+
"""Compute the length of a floating-point vector ``a``."""
|
|
464
|
+
...
|
|
465
|
+
|
|
466
|
+
@over
|
|
467
|
+
def length(a: Quaternion[Float]) -> Float:
|
|
468
|
+
"""Compute the length of a quaternion ``a``."""
|
|
469
|
+
...
|
|
470
|
+
|
|
471
|
+
@over
|
|
472
|
+
def length_sq(a: Vector[Any, Scalar]) -> Scalar:
|
|
473
|
+
"""Compute the squared length of a vector ``a``."""
|
|
474
|
+
...
|
|
475
|
+
|
|
476
|
+
@over
|
|
477
|
+
def length_sq(a: Quaternion[Scalar]) -> Scalar:
|
|
478
|
+
"""Compute the squared length of a quaternion ``a``."""
|
|
479
|
+
...
|
|
480
|
+
|
|
481
|
+
@over
|
|
482
|
+
def normalize(a: Vector[Any, Float]) -> Vector[Any, Float]:
|
|
483
|
+
"""Compute the normalized value of ``a``. If ``length(a)`` is 0 then the zero vector is returned."""
|
|
484
|
+
...
|
|
485
|
+
|
|
486
|
+
@over
|
|
487
|
+
def normalize(a: Quaternion[Float]) -> Quaternion[Float]:
|
|
488
|
+
"""Compute the normalized value of ``a``. If ``length(a)`` is 0, then the zero quaternion is returned."""
|
|
489
|
+
...
|
|
490
|
+
|
|
491
|
+
@over
|
|
492
|
+
def transpose(a: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
493
|
+
"""Return the transpose of the matrix ``a``."""
|
|
494
|
+
...
|
|
495
|
+
|
|
496
|
+
@over
|
|
497
|
+
def inverse(a: Matrix[2, 2, Float]) -> Matrix[Any, Any, Float]:
|
|
498
|
+
"""Return the inverse of a 2x2 matrix ``a``."""
|
|
499
|
+
...
|
|
500
|
+
|
|
501
|
+
@over
|
|
502
|
+
def inverse(a: Matrix[3, 3, Float]) -> Matrix[Any, Any, Float]:
|
|
503
|
+
"""Return the inverse of a 3x3 matrix ``a``."""
|
|
504
|
+
...
|
|
505
|
+
|
|
506
|
+
@over
|
|
507
|
+
def inverse(a: Matrix[4, 4, Float]) -> Matrix[Any, Any, Float]:
|
|
508
|
+
"""Return the inverse of a 4x4 matrix ``a``."""
|
|
509
|
+
...
|
|
510
|
+
|
|
511
|
+
@over
|
|
512
|
+
def determinant(a: Matrix[2, 2, Float]) -> Float:
|
|
513
|
+
"""Return the determinant of a 2x2 matrix ``a``."""
|
|
514
|
+
...
|
|
515
|
+
|
|
516
|
+
@over
|
|
517
|
+
def determinant(a: Matrix[3, 3, Float]) -> Float:
|
|
518
|
+
"""Return the determinant of a 3x3 matrix ``a``."""
|
|
519
|
+
...
|
|
520
|
+
|
|
521
|
+
@over
|
|
522
|
+
def determinant(a: Matrix[4, 4, Float]) -> Float:
|
|
523
|
+
"""Return the determinant of a 4x4 matrix ``a``."""
|
|
524
|
+
...
|
|
525
|
+
|
|
526
|
+
@over
|
|
527
|
+
def trace(a: Matrix[Any, Any, Scalar]) -> Scalar:
|
|
528
|
+
"""Return the trace of the matrix ``a``."""
|
|
529
|
+
...
|
|
530
|
+
|
|
531
|
+
@over
|
|
532
|
+
def diag(vec: Vector[Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
533
|
+
"""Returns a matrix with the components of the vector ``vec`` on the diagonal."""
|
|
534
|
+
...
|
|
535
|
+
|
|
536
|
+
@over
|
|
537
|
+
def get_diag(mat: Matrix[Any, Any, Scalar]) -> Vector[Any, Scalar]:
|
|
538
|
+
"""Returns a vector containing the diagonal elements of the square matrix ``mat``."""
|
|
539
|
+
...
|
|
540
|
+
|
|
541
|
+
@over
|
|
542
|
+
def cw_mul(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
543
|
+
"""Component-wise multiplication of two vectors."""
|
|
544
|
+
...
|
|
545
|
+
|
|
546
|
+
@over
|
|
547
|
+
def cw_mul(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
548
|
+
"""Component-wise multiplication of two matrices."""
|
|
549
|
+
...
|
|
550
|
+
|
|
551
|
+
@over
|
|
552
|
+
def cw_div(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
553
|
+
"""Component-wise division of two vectors."""
|
|
554
|
+
...
|
|
555
|
+
|
|
556
|
+
@over
|
|
557
|
+
def cw_div(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
558
|
+
"""Component-wise division of two matrices."""
|
|
559
|
+
...
|
|
560
|
+
|
|
561
|
+
@over
|
|
562
|
+
def vector(*args: Scalar, length: int32, dtype: Scalar) -> Vector[Any, Scalar]:
|
|
563
|
+
"""Construct a vector of given length and dtype."""
|
|
564
|
+
...
|
|
565
|
+
|
|
566
|
+
@over
|
|
567
|
+
def matrix(pos: Vector[3, Float], rot: Quaternion[Float], scale: Vector[3, Float], dtype: Float) -> Matrix[4, 4, Float]:
|
|
568
|
+
"""Construct a 4x4 transformation matrix that applies the transformations as
|
|
569
|
+
Translation(pos)*Rotation(rot)*Scaling(scale) when applied to column vectors, i.e.: y = (TRS)*x
|
|
570
|
+
|
|
571
|
+
.. warning::
|
|
572
|
+
This function has been deprecated in favor of :func:`warp.math.transform_compose()`.
|
|
573
|
+
"""
|
|
574
|
+
...
|
|
575
|
+
|
|
576
|
+
@over
|
|
577
|
+
def matrix(*args: Scalar, shape: Tuple[int, int], dtype: Scalar) -> Matrix[Any, Any, Scalar]:
|
|
578
|
+
"""Construct a matrix. If the positional ``arg_types`` are not given, then matrix will be zero-initialized."""
|
|
579
|
+
...
|
|
580
|
+
|
|
581
|
+
@over
|
|
582
|
+
def matrix_from_cols(*args: Vector[Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
583
|
+
"""Construct a matrix from column vectors."""
|
|
584
|
+
...
|
|
585
|
+
|
|
586
|
+
@over
|
|
587
|
+
def matrix_from_rows(*args: Vector[Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
588
|
+
"""Construct a matrix from row vectors."""
|
|
589
|
+
...
|
|
590
|
+
|
|
591
|
+
@over
|
|
592
|
+
def identity(n: int32, dtype: Scalar) -> Matrix[Any, Any, Scalar]:
|
|
593
|
+
"""Create an identity matrix with shape=(n,n) with the type given by ``dtype``."""
|
|
594
|
+
...
|
|
595
|
+
|
|
596
|
+
@over
|
|
597
|
+
def svd3(A: Matrix[3, 3, Float]) -> Tuple[Matrix[3, 3, Float], Vector[3, Float], Matrix[3, 3, Float]]:
|
|
598
|
+
"""Compute the SVD of a 3x3 matrix ``A``. The singular values are returned in ``sigma``,
|
|
599
|
+
while the left and right basis vectors are returned in ``U`` and ``V``.
|
|
600
|
+
"""
|
|
601
|
+
...
|
|
602
|
+
|
|
603
|
+
@over
|
|
604
|
+
def svd3(A: Matrix[3, 3, Float], U: Matrix[3, 3, Float], sigma: Vector[3, Float], V: Matrix[3, 3, Float]):
|
|
605
|
+
"""Compute the SVD of a 3x3 matrix ``A``. The singular values are returned in ``sigma``,
|
|
606
|
+
while the left and right basis vectors are returned in ``U`` and ``V``.
|
|
607
|
+
"""
|
|
608
|
+
...
|
|
609
|
+
|
|
610
|
+
@over
|
|
611
|
+
def svd2(A: Matrix[2, 2, Float]) -> Tuple[Matrix[2, 2, Float], Vector[2, Float], Matrix[2, 2, Float]]:
|
|
612
|
+
"""Compute the SVD of a 2x2 matrix ``A``. The singular values are returned in ``sigma``,
|
|
613
|
+
while the left and right basis vectors are returned in ``U`` and ``V``.
|
|
614
|
+
"""
|
|
615
|
+
...
|
|
616
|
+
|
|
617
|
+
@over
|
|
618
|
+
def svd2(A: Matrix[2, 2, Float], U: Matrix[2, 2, Float], sigma: Vector[2, Float], V: Matrix[2, 2, Float]):
|
|
619
|
+
"""Compute the SVD of a 2x2 matrix ``A``. The singular values are returned in ``sigma``,
|
|
620
|
+
while the left and right basis vectors are returned in ``U`` and ``V``.
|
|
621
|
+
"""
|
|
622
|
+
...
|
|
623
|
+
|
|
624
|
+
@over
|
|
625
|
+
def qr3(A: Matrix[3, 3, Float]) -> Tuple[Matrix[3, 3, Float], Matrix[3, 3, Float]]:
|
|
626
|
+
"""Compute the QR decomposition of a 3x3 matrix ``A``. The orthogonal matrix is returned in ``Q``,
|
|
627
|
+
while the upper triangular matrix is returned in ``R``.
|
|
628
|
+
"""
|
|
629
|
+
...
|
|
630
|
+
|
|
631
|
+
@over
|
|
632
|
+
def qr3(A: Matrix[3, 3, Float], Q: Matrix[3, 3, Float], R: Matrix[3, 3, Float]):
|
|
633
|
+
"""Compute the QR decomposition of a 3x3 matrix ``A``. The orthogonal matrix is returned in ``Q``,
|
|
634
|
+
while the upper triangular matrix is returned in ``R``.
|
|
635
|
+
"""
|
|
636
|
+
...
|
|
637
|
+
|
|
638
|
+
@over
|
|
639
|
+
def eig3(A: Matrix[3, 3, Float]) -> Tuple[Matrix[3, 3, Float], Vector[3, Float]]:
|
|
640
|
+
"""Compute the eigendecomposition of a 3x3 matrix ``A``. The eigenvectors are returned as the columns of ``Q``,
|
|
641
|
+
while the corresponding eigenvalues are returned in ``d``.
|
|
642
|
+
"""
|
|
643
|
+
...
|
|
644
|
+
|
|
645
|
+
@over
|
|
646
|
+
def eig3(A: Matrix[3, 3, Float], Q: Matrix[3, 3, Float], d: Vector[3, Float]):
|
|
647
|
+
"""Compute the eigendecomposition of a 3x3 matrix ``A``. The eigenvectors are returned as the columns of ``Q``,
|
|
648
|
+
while the corresponding eigenvalues are returned in ``d``.
|
|
649
|
+
"""
|
|
650
|
+
...
|
|
651
|
+
|
|
652
|
+
@over
|
|
653
|
+
def quaternion(dtype: Float) -> Quaternion[Float]:
|
|
654
|
+
"""Construct a zero-initialized quaternion. Quaternions are laid out as
|
|
655
|
+
[ix, iy, iz, r], where ix, iy, iz are the imaginary part, and r the real part.
|
|
656
|
+
"""
|
|
657
|
+
...
|
|
658
|
+
|
|
659
|
+
@over
|
|
660
|
+
def quaternion(x: Float, y: Float, z: Float, w: Float, dtype: Scalar) -> Quaternion[Float]:
|
|
661
|
+
"""Create a quaternion using the supplied components (type inferred from component type)."""
|
|
662
|
+
...
|
|
663
|
+
|
|
664
|
+
@over
|
|
665
|
+
def quaternion(ijk: Vector[3, Float], real: Float, dtype: Float) -> Quaternion[Float]:
|
|
666
|
+
"""Create a quaternion using the supplied vector/scalar (type inferred from scalar type)."""
|
|
667
|
+
...
|
|
668
|
+
|
|
669
|
+
@over
|
|
670
|
+
def quaternion(quat: Quaternion[Float], dtype: Float) -> Quaternion[Float]:
|
|
671
|
+
"""Construct a quaternion of type dtype from another quaternion of a different dtype."""
|
|
672
|
+
...
|
|
673
|
+
|
|
674
|
+
@over
|
|
675
|
+
def quat_identity(dtype: Float) -> quatf:
|
|
676
|
+
"""Construct an identity quaternion with zero imaginary part and real part of 1.0"""
|
|
677
|
+
...
|
|
678
|
+
|
|
679
|
+
@over
|
|
680
|
+
def quat_from_axis_angle(axis: Vector[3, Float], angle: Float) -> Quaternion[Float]:
|
|
681
|
+
"""Construct a quaternion representing a rotation of angle radians around the given axis."""
|
|
682
|
+
...
|
|
683
|
+
|
|
684
|
+
@over
|
|
685
|
+
def quat_to_axis_angle(quat: Quaternion[Float]) -> Tuple[Vector[3, Float], Float]:
|
|
686
|
+
"""Extract the rotation axis and angle radians a quaternion represents."""
|
|
687
|
+
...
|
|
688
|
+
|
|
689
|
+
@over
|
|
690
|
+
def quat_to_axis_angle(quat: Quaternion[Float], axis: Vector[3, Float], angle: Float):
|
|
691
|
+
"""Extract the rotation axis and angle radians a quaternion represents."""
|
|
692
|
+
...
|
|
693
|
+
|
|
694
|
+
@over
|
|
695
|
+
def quat_from_matrix(mat: Matrix[3, 3, Float]) -> Quaternion[Float]:
|
|
696
|
+
"""Construct a quaternion from a 3x3 matrix.
|
|
697
|
+
|
|
698
|
+
If the matrix is not a pure rotation, but for example includes scaling or skewing, the result is undefined.
|
|
699
|
+
"""
|
|
700
|
+
...
|
|
701
|
+
|
|
702
|
+
@over
|
|
703
|
+
def quat_from_matrix(mat: Matrix[4, 4, Float]) -> Quaternion[Float]:
|
|
704
|
+
"""Construct a quaternion from a 4x4 matrix.
|
|
705
|
+
|
|
706
|
+
If the top-left 3x3 block of the matrix is not a pure rotation, but for example includes scaling or skewing, the result is undefined.
|
|
707
|
+
"""
|
|
708
|
+
...
|
|
709
|
+
|
|
710
|
+
@over
|
|
711
|
+
def quat_rpy(roll: Float, pitch: Float, yaw: Float) -> Quaternion[Float]:
|
|
712
|
+
"""Construct a quaternion representing a combined roll (z), pitch (x), yaw rotations (y) in radians."""
|
|
713
|
+
...
|
|
714
|
+
|
|
715
|
+
@over
|
|
716
|
+
def quat_inverse(quat: Quaternion[Float]) -> Quaternion[Float]:
|
|
717
|
+
"""Compute quaternion conjugate."""
|
|
718
|
+
...
|
|
719
|
+
|
|
720
|
+
@over
|
|
721
|
+
def quat_rotate(quat: Quaternion[Float], vec: Vector[3, Float]) -> Vector[3, Float]:
|
|
722
|
+
"""Rotate a vector by a quaternion."""
|
|
723
|
+
...
|
|
724
|
+
|
|
725
|
+
@over
|
|
726
|
+
def quat_rotate_inv(quat: Quaternion[Float], vec: Vector[3, Float]) -> Vector[3, Float]:
|
|
727
|
+
"""Rotate a vector by the inverse of a quaternion."""
|
|
728
|
+
...
|
|
729
|
+
|
|
730
|
+
@over
|
|
731
|
+
def quat_slerp(a: Quaternion[Float], b: Quaternion[Float], t: Float) -> Quaternion[Float]:
|
|
732
|
+
"""Linearly interpolate between two quaternions."""
|
|
733
|
+
...
|
|
734
|
+
|
|
735
|
+
@over
|
|
736
|
+
def quat_to_matrix(quat: Quaternion[Float]) -> Matrix[3, 3, Float]:
|
|
737
|
+
"""Convert a quaternion to a 3x3 rotation matrix."""
|
|
738
|
+
...
|
|
739
|
+
|
|
740
|
+
@over
|
|
741
|
+
def transformation(p: Vector[3, Float], q: Quaternion[Float], dtype: Float) -> Transformation[Float]:
|
|
742
|
+
"""Construct a rigid-body transformation with translation part ``p`` and rotation ``q``."""
|
|
743
|
+
...
|
|
744
|
+
|
|
745
|
+
@over
|
|
746
|
+
def transformation(*args: Float, dtype: Float) -> Transformation[Float]:
|
|
747
|
+
"""Construct a spatial transform vector of given dtype."""
|
|
748
|
+
...
|
|
749
|
+
|
|
750
|
+
@over
|
|
751
|
+
def transform_identity(dtype: Float) -> transformf:
|
|
752
|
+
"""Construct an identity transform with zero translation and identity rotation."""
|
|
753
|
+
...
|
|
754
|
+
|
|
755
|
+
@over
|
|
756
|
+
def transform_get_translation(xform: Transformation[Float]) -> Vector[3, Float]:
|
|
757
|
+
"""Return the translational part of a transform ``xform``."""
|
|
758
|
+
...
|
|
759
|
+
|
|
760
|
+
@over
|
|
761
|
+
def transform_get_rotation(xform: Transformation[Float]) -> Quaternion[Float]:
|
|
762
|
+
"""Return the rotational part of a transform ``xform``."""
|
|
763
|
+
...
|
|
764
|
+
|
|
765
|
+
@over
|
|
766
|
+
def transform_set_translation(xform: Transformation[Float], p: Vector[3, Float]):
|
|
767
|
+
"""Set the translational part of a transform ``xform``."""
|
|
768
|
+
...
|
|
769
|
+
|
|
770
|
+
@over
|
|
771
|
+
def transform_set_rotation(xform: Transformation[Float], q: Quaternion[Float]):
|
|
772
|
+
"""Set the rotational part of a transform ``xform``."""
|
|
773
|
+
...
|
|
774
|
+
|
|
775
|
+
@over
|
|
776
|
+
def transform_multiply(a: Transformation[Float], b: Transformation[Float]) -> Transformation[Float]:
|
|
777
|
+
"""Multiply two rigid body transformations together."""
|
|
778
|
+
...
|
|
779
|
+
|
|
780
|
+
@over
|
|
781
|
+
def transform_point(xform: Transformation[Float], point: Vector[3, Float]) -> Vector[3, Float]:
|
|
782
|
+
"""Apply the transform to a point ``point`` treating the homogeneous coordinate as w=1 (translation and rotation)."""
|
|
783
|
+
...
|
|
784
|
+
|
|
785
|
+
@over
|
|
786
|
+
def transform_point(mat: Matrix[4, 4, Float], point: Vector[3, Float]) -> Vector[3, Float]:
|
|
787
|
+
"""Apply the transform to a point ``point`` treating the homogeneous coordinate as w=1.
|
|
788
|
+
|
|
789
|
+
The transformation is applied treating ``point`` as a column vector, e.g.: ``y = mat*point``.
|
|
790
|
+
|
|
791
|
+
This is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = point^T*mat^T``.
|
|
792
|
+
If the transform is coming from a library that uses row-vectors, then users should transpose the transformation
|
|
793
|
+
matrix before calling this method.
|
|
794
|
+
"""
|
|
795
|
+
...
|
|
796
|
+
|
|
797
|
+
@over
|
|
798
|
+
def transform_vector(xform: Transformation[Float], vec: Vector[3, Float]) -> Vector[3, Float]:
|
|
799
|
+
"""Apply the transform to a vector ``vec`` treating the homogeneous coordinate as w=0 (rotation only)."""
|
|
800
|
+
...
|
|
801
|
+
|
|
802
|
+
@over
|
|
803
|
+
def transform_vector(mat: Matrix[4, 4, Float], vec: Vector[3, Float]) -> Vector[3, Float]:
|
|
804
|
+
"""Apply the transform to a vector ``vec`` treating the homogeneous coordinate as w=0.
|
|
805
|
+
|
|
806
|
+
The transformation is applied treating ``vec`` as a column vector, e.g.: ``y = mat*vec``.
|
|
807
|
+
|
|
808
|
+
This is in contrast to some libraries, notably USD, which applies transforms to row vectors, ``y^T = vec^T*mat^T``.
|
|
809
|
+
If the transform is coming from a library that uses row-vectors, then users should transpose the transformation
|
|
810
|
+
matrix before calling this method.
|
|
811
|
+
"""
|
|
812
|
+
...
|
|
813
|
+
|
|
814
|
+
@over
|
|
815
|
+
def transform_inverse(xform: Transformation[Float]) -> Transformation[Float]:
|
|
816
|
+
"""Compute the inverse of the transformation ``xform``."""
|
|
817
|
+
...
|
|
818
|
+
|
|
819
|
+
@over
|
|
820
|
+
def spatial_vector(dtype: Float) -> Vector[6, Float]:
|
|
821
|
+
"""Zero-initialize a 6D screw vector."""
|
|
822
|
+
...
|
|
823
|
+
|
|
824
|
+
@over
|
|
825
|
+
def spatial_vector(w: Vector[3, Float], v: Vector[3, Float], dtype: Float) -> Vector[6, Float]:
|
|
826
|
+
"""Construct a 6D screw vector from two 3D vectors."""
|
|
827
|
+
...
|
|
828
|
+
|
|
829
|
+
@over
|
|
830
|
+
def spatial_vector(wx: Float, wy: Float, wz: Float, vx: Float, vy: Float, vz: Float, dtype: Float) -> Vector[6, Float]:
|
|
831
|
+
"""Construct a 6D screw vector from six values."""
|
|
832
|
+
...
|
|
833
|
+
|
|
834
|
+
@over
|
|
835
|
+
def spatial_adjoint(r: Matrix[3, 3, Float], s: Matrix[3, 3, Float]) -> Matrix[6, 6, Float]:
|
|
836
|
+
"""Construct a 6x6 spatial inertial matrix from two 3x3 diagonal blocks."""
|
|
837
|
+
...
|
|
838
|
+
|
|
839
|
+
@over
|
|
840
|
+
def spatial_dot(a: Vector[6, Float], b: Vector[6, Float]) -> Float:
|
|
841
|
+
"""Compute the dot product of two 6D screw vectors."""
|
|
842
|
+
...
|
|
843
|
+
|
|
844
|
+
@over
|
|
845
|
+
def spatial_cross(a: Vector[6, Float], b: Vector[6, Float]) -> Vector[6, Float]:
|
|
846
|
+
"""Compute the cross product of two 6D screw vectors."""
|
|
847
|
+
...
|
|
848
|
+
|
|
849
|
+
@over
|
|
850
|
+
def spatial_cross_dual(a: Vector[6, Float], b: Vector[6, Float]) -> Vector[6, Float]:
|
|
851
|
+
"""Compute the dual cross product of two 6D screw vectors."""
|
|
852
|
+
...
|
|
853
|
+
|
|
854
|
+
@over
|
|
855
|
+
def spatial_top(svec: Vector[6, Float]) -> Vector[3, Float]:
|
|
856
|
+
"""Return the top (first) part of a 6D screw vector."""
|
|
857
|
+
...
|
|
858
|
+
|
|
859
|
+
@over
|
|
860
|
+
def spatial_bottom(svec: Vector[6, Float]) -> Vector[3, Float]:
|
|
861
|
+
"""Return the bottom (second) part of a 6D screw vector."""
|
|
862
|
+
...
|
|
863
|
+
|
|
864
|
+
@over
|
|
865
|
+
def spatial_jacobian(
|
|
866
|
+
S: Array[Vector[6, Float]],
|
|
867
|
+
joint_parents: Array[int32],
|
|
868
|
+
joint_qd_start: Array[int32],
|
|
869
|
+
joint_start: int32,
|
|
870
|
+
joint_count: int32,
|
|
871
|
+
J_start: int32,
|
|
872
|
+
J_out: Array[Float],
|
|
873
|
+
):
|
|
874
|
+
""" """
|
|
875
|
+
...
|
|
876
|
+
|
|
877
|
+
@over
|
|
878
|
+
def spatial_mass(
|
|
879
|
+
I_s: Array[Matrix[6, 6, Float]], joint_start: int32, joint_count: int32, M_start: int32, M: Array[Float]
|
|
880
|
+
):
|
|
881
|
+
""" """
|
|
882
|
+
...
|
|
883
|
+
|
|
884
|
+
@over
|
|
885
|
+
def tile_zeros(shape: Tuple[int, ...], dtype: Any, storage: str) -> Tile[Any, Tuple[int, ...]]:
|
|
886
|
+
"""Allocate a tile of zero-initialized items.
|
|
887
|
+
|
|
888
|
+
:param shape: Shape of the output tile
|
|
889
|
+
:param dtype: Data type of output tile's elements (default float)
|
|
890
|
+
:param storage: The storage location for the tile: ``"register"`` for registers
|
|
891
|
+
(default) or ``"shared"`` for shared memory.
|
|
892
|
+
:returns: A zero-initialized tile with shape and data type as specified
|
|
893
|
+
"""
|
|
894
|
+
...
|
|
895
|
+
|
|
896
|
+
@over
|
|
897
|
+
def tile_ones(shape: Tuple[int, ...], dtype: Any, storage: str) -> Tile[Any, Tuple[int, ...]]:
|
|
898
|
+
"""Allocate a tile of one-initialized items.
|
|
899
|
+
|
|
900
|
+
:param shape: Shape of the output tile
|
|
901
|
+
:param dtype: Data type of output tile's elements
|
|
902
|
+
:param storage: The storage location for the tile: ``"register"`` for registers
|
|
903
|
+
(default) or ``"shared"`` for shared memory.
|
|
904
|
+
:returns: A one-initialized tile with shape and data type as specified
|
|
905
|
+
"""
|
|
906
|
+
...
|
|
907
|
+
|
|
908
|
+
@over
|
|
909
|
+
def tile_arange(*args: Scalar, dtype: Scalar, storage: str) -> Tile[Scalar, Tuple[int]]:
|
|
910
|
+
"""Generate a tile of linearly spaced elements.
|
|
911
|
+
|
|
912
|
+
:param args: Variable-length positional arguments, interpreted as:
|
|
913
|
+
|
|
914
|
+
- ``(stop,)``: Generates values from ``0`` to ``stop - 1``
|
|
915
|
+
- ``(start, stop)``: Generates values from ``start`` to ``stop - 1``
|
|
916
|
+
- ``(start, stop, step)``: Generates values from ``start`` to ``stop - 1`` with a step size
|
|
917
|
+
|
|
918
|
+
:param dtype: Data type of output tile's elements (optional, default: ``float``)
|
|
919
|
+
:param storage: The storage location for the tile: ``"register"`` for registers
|
|
920
|
+
(default) or ``"shared"`` for shared memory.
|
|
921
|
+
:returns: A tile with ``shape=(n)`` with linearly spaced elements of specified data type
|
|
922
|
+
"""
|
|
923
|
+
...
|
|
924
|
+
|
|
925
|
+
@over
|
|
926
|
+
def tile_load(
|
|
927
|
+
a: Array[Any], shape: Tuple[int, ...], offset: Tuple[int, ...], storage: str
|
|
928
|
+
) -> Tile[Any, Tuple[int, ...]]:
|
|
929
|
+
"""Loads a tile from a global memory array.
|
|
930
|
+
|
|
931
|
+
This method will cooperatively load a tile from global memory using all threads in the block.
|
|
932
|
+
|
|
933
|
+
:param a: The source array in global memory
|
|
934
|
+
:param shape: Shape of the tile to load, must have the same number of dimensions as ``a``
|
|
935
|
+
:param offset: Offset in the source array to begin reading from (optional)
|
|
936
|
+
:param storage: The storage location for the tile: ``"register"`` for registers
|
|
937
|
+
(default) or ``"shared"`` for shared memory.
|
|
938
|
+
:returns: A tile with shape as specified and data type the same as the source array
|
|
939
|
+
"""
|
|
940
|
+
...
|
|
941
|
+
|
|
942
|
+
@over
|
|
943
|
+
def tile_store(a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...]):
|
|
944
|
+
"""Store a tile to a global memory array.
|
|
945
|
+
|
|
946
|
+
This method will cooperatively store a tile to global memory using all threads in the block.
|
|
947
|
+
|
|
948
|
+
:param a: The destination array in global memory
|
|
949
|
+
:param t: The source tile to store data from, must have the same data type and number of dimensions as the destination array
|
|
950
|
+
:param offset: Offset in the destination array (optional)
|
|
951
|
+
"""
|
|
952
|
+
...
|
|
953
|
+
|
|
954
|
+
@over
|
|
955
|
+
def tile_atomic_add(
|
|
956
|
+
a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...]
|
|
957
|
+
) -> Tile[Any, Tuple[int, ...]]:
|
|
958
|
+
"""Atomically add a tile onto the array `a`, each element will be updated atomically.
|
|
959
|
+
|
|
960
|
+
:param a: Array in global memory, should have the same ``dtype`` as the input tile
|
|
961
|
+
:param t: Source tile to add to the destination array
|
|
962
|
+
:param offset: Offset in the destination array (optional)
|
|
963
|
+
:returns: A tile with the same dimensions and data type as the source tile, holding the original value of the destination elements
|
|
964
|
+
"""
|
|
965
|
+
...
|
|
966
|
+
|
|
967
|
+
@over
|
|
968
|
+
def tile_view(
|
|
969
|
+
t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...], shape: Tuple[int, ...]
|
|
970
|
+
) -> Tile[Any, Tuple[int, ...]]:
|
|
971
|
+
"""Return a slice of a given tile [offset, offset+shape], if shape is not specified it will be inferred from the unspecified offset dimensions.
|
|
972
|
+
|
|
973
|
+
:param t: Input tile to extract a subrange from
|
|
974
|
+
:param offset: Offset in the source tile
|
|
975
|
+
:param shape: Shape of the returned slice
|
|
976
|
+
:returns: A tile with dimensions given by the specified shape or the remaining source tile dimensions
|
|
977
|
+
"""
|
|
978
|
+
...
|
|
979
|
+
|
|
980
|
+
@over
|
|
981
|
+
def tile_squeeze(t: Tile[Any, Tuple[int, ...]], axis: Tuple[int, ...]) -> Tile[Any, Tuple[int, ...]]:
|
|
982
|
+
"""Return a squeezed view of a tile with the same data.
|
|
983
|
+
|
|
984
|
+
:param t: Input tile to squeeze
|
|
985
|
+
:param axis: A subset of the entries of length one in the shape (optional)
|
|
986
|
+
:returns: The input tile but with all or a subset of the dimensions of length one removed.
|
|
987
|
+
"""
|
|
988
|
+
...
|
|
989
|
+
|
|
990
|
+
@over
|
|
991
|
+
def tile_reshape(t: Tile[Any, Tuple[int, ...]], shape: Tuple[int, ...]) -> Tile[Any, Tuple[int, ...]]:
|
|
992
|
+
"""Return a reshaped view of a tile with the same data.
|
|
993
|
+
|
|
994
|
+
:param t: Input tile to reshape
|
|
995
|
+
:param shape: New shape for the tile
|
|
996
|
+
:returns: A tile containing the same data as the input tile, but arranged in a new shape.
|
|
997
|
+
"""
|
|
998
|
+
...
|
|
999
|
+
|
|
1000
|
+
@over
|
|
1001
|
+
def tile_astype(t: Tile[Scalar, Tuple[int, ...]], dtype: Scalar) -> Tile[Any, Tuple[int, ...]]:
|
|
1002
|
+
"""Return a new tile with the same data as the input tile, but with a different data type.
|
|
1003
|
+
|
|
1004
|
+
:param t: Input tile
|
|
1005
|
+
:param dtype: New data type for the tile
|
|
1006
|
+
:returns: A tile with the same data as the input tile, but with a different data type
|
|
1007
|
+
"""
|
|
1008
|
+
...
|
|
1009
|
+
|
|
1010
|
+
@over
|
|
1011
|
+
def tile_assign(dst: Tile[Any, Tuple[int, ...]], src: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...]):
|
|
1012
|
+
"""Assign a tile to a subrange of a destination tile.
|
|
1013
|
+
|
|
1014
|
+
:param dst: The destination tile to assign to
|
|
1015
|
+
:param src: The source tile to read values from
|
|
1016
|
+
:param offset: Offset in the destination tile to write to
|
|
1017
|
+
"""
|
|
1018
|
+
...
|
|
1019
|
+
|
|
1020
|
+
@over
|
|
1021
|
+
def untile(a: Tile[Any, Tuple[int, ...]]) -> Any:
|
|
1022
|
+
"""Convert a tile back to per-thread values.
|
|
1023
|
+
|
|
1024
|
+
This function converts a block-wide tile back to per-thread values.
|
|
1025
|
+
|
|
1026
|
+
* If the input tile is 1D, then the resulting value will be a per-thread scalar
|
|
1027
|
+
* If the input tile is 2D, then the resulting value will be a per-thread vector of length M
|
|
1028
|
+
|
|
1029
|
+
:param a: A tile with dimensions ``shape=(M, block_dim)``
|
|
1030
|
+
:returns: A single value per-thread with the same data type as the tile
|
|
1031
|
+
|
|
1032
|
+
This example shows how to create a linear sequence from thread variables:
|
|
1033
|
+
|
|
1034
|
+
.. code-block:: python
|
|
1035
|
+
|
|
1036
|
+
@wp.kernel
|
|
1037
|
+
def compute():
|
|
1038
|
+
i = wp.tid()
|
|
1039
|
+
|
|
1040
|
+
# create block-wide tile
|
|
1041
|
+
t = wp.tile(i) * 2
|
|
1042
|
+
|
|
1043
|
+
# convert back to per-thread values
|
|
1044
|
+
s = wp.untile(t)
|
|
1045
|
+
|
|
1046
|
+
print(s)
|
|
1047
|
+
|
|
1048
|
+
wp.launch(compute, dim=16, inputs=[], block_dim=16)
|
|
1049
|
+
|
|
1050
|
+
Prints:
|
|
1051
|
+
|
|
1052
|
+
.. code-block:: text
|
|
1053
|
+
|
|
1054
|
+
0
|
|
1055
|
+
2
|
|
1056
|
+
4
|
|
1057
|
+
6
|
|
1058
|
+
8
|
|
1059
|
+
...
|
|
1060
|
+
|
|
1061
|
+
"""
|
|
1062
|
+
...
|
|
1063
|
+
|
|
1064
|
+
@over
|
|
1065
|
+
def tile_transpose(a: Tile[Any, Tuple[int, int]]) -> Tile[Any, Tuple[int, int]]:
|
|
1066
|
+
"""Transpose a tile.
|
|
1067
|
+
|
|
1068
|
+
For shared memory tiles, this operation will alias the input tile.
|
|
1069
|
+
Register tiles will first be transferred to shared memory before transposition.
|
|
1070
|
+
|
|
1071
|
+
:param a: Tile to transpose with ``shape=(M,N)``
|
|
1072
|
+
:returns: Tile with ``shape=(N,M)``
|
|
1073
|
+
"""
|
|
1074
|
+
...
|
|
1075
|
+
|
|
1076
|
+
@over
|
|
1077
|
+
def tile_broadcast(a: Tile[Any, Tuple[int, ...]], shape: Tuple[int, ...]) -> Tile[Any, Tuple[int, ...]]:
|
|
1078
|
+
"""Broadcast a tile.
|
|
1079
|
+
|
|
1080
|
+
Broadcasts the input tile ``a`` to the destination shape.
|
|
1081
|
+
Broadcasting follows NumPy broadcast rules.
|
|
1082
|
+
|
|
1083
|
+
:param a: Tile to broadcast
|
|
1084
|
+
:param shape: The shape to broadcast to
|
|
1085
|
+
:returns: Tile with broadcast shape
|
|
1086
|
+
"""
|
|
1087
|
+
...
|
|
1088
|
+
|
|
1089
|
+
@over
|
|
1090
|
+
def tile_sum(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[1]]:
|
|
1091
|
+
"""Cooperatively compute the sum of the tile elements using all threads in the block.
|
|
1092
|
+
|
|
1093
|
+
:param a: The tile to compute the sum of
|
|
1094
|
+
:returns: A single-element tile holding the sum
|
|
1095
|
+
|
|
1096
|
+
Example:
|
|
1097
|
+
|
|
1098
|
+
.. code-block:: python
|
|
1099
|
+
|
|
1100
|
+
@wp.kernel
|
|
1101
|
+
def compute():
|
|
1102
|
+
t = wp.tile_ones(dtype=float, shape=(16, 16))
|
|
1103
|
+
s = wp.tile_sum(t)
|
|
1104
|
+
|
|
1105
|
+
print(s)
|
|
1106
|
+
|
|
1107
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
|
|
1108
|
+
|
|
1109
|
+
Prints:
|
|
1110
|
+
|
|
1111
|
+
.. code-block:: text
|
|
1112
|
+
|
|
1113
|
+
[256] = tile(shape=(1), storage=register)
|
|
1114
|
+
|
|
1115
|
+
|
|
1116
|
+
"""
|
|
1117
|
+
...
|
|
1118
|
+
|
|
1119
|
+
@over
|
|
1120
|
+
def tile_sort(keys: Tile[Any, Tuple[int]], values: Tile[Any, Tuple[int]]):
|
|
1121
|
+
"""Cooperatively sort the elements of two tiles in ascending order based on the keys, using all threads in the block.
|
|
1122
|
+
|
|
1123
|
+
:param keys: Keys to sort by. Supported key types: :class:`float32`, :class:`int32`, :class:`uint32`. Must be in shared memory.
|
|
1124
|
+
:param values: Values to sort along with keys. No type restrictions. Must be in shared memory.
|
|
1125
|
+
:returns: No return value. Sorts both tiles in-place.
|
|
1126
|
+
|
|
1127
|
+
Example:
|
|
1128
|
+
|
|
1129
|
+
.. code-block:: python
|
|
1130
|
+
|
|
1131
|
+
@wp.kernel
|
|
1132
|
+
def compute():
|
|
1133
|
+
keys = wp.tile_arange(32, 0, -1, dtype=int, storage="shared")
|
|
1134
|
+
values = wp.tile_arange(0, 32, 1, dtype=int, storage="shared")
|
|
1135
|
+
wp.tile_sort(keys, values)
|
|
1136
|
+
|
|
1137
|
+
print(keys)
|
|
1138
|
+
print(values)
|
|
1139
|
+
|
|
1140
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
|
|
1141
|
+
|
|
1142
|
+
Prints:
|
|
1143
|
+
|
|
1144
|
+
.. code-block:: text
|
|
1145
|
+
|
|
1146
|
+
[1, 2, ..., 32] = tile(shape=(32), storage=shared)
|
|
1147
|
+
[31, 30, 29, ..., 0] = tile(shape=(32), storage=shared)
|
|
1148
|
+
|
|
1149
|
+
|
|
1150
|
+
"""
|
|
1151
|
+
...
|
|
1152
|
+
|
|
1153
|
+
@over
|
|
1154
|
+
def tile_min(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[1]]:
|
|
1155
|
+
"""Cooperatively compute the minimum of the tile elements using all threads in the block.
|
|
1156
|
+
|
|
1157
|
+
:param a: The tile to compute the minimum of
|
|
1158
|
+
:returns: A single-element tile holding the minimum value
|
|
1159
|
+
|
|
1160
|
+
Example:
|
|
1161
|
+
|
|
1162
|
+
.. code-block:: python
|
|
1163
|
+
|
|
1164
|
+
@wp.kernel
|
|
1165
|
+
def compute():
|
|
1166
|
+
t = wp.tile_arange(64, 128)
|
|
1167
|
+
s = wp.tile_min(t)
|
|
1168
|
+
|
|
1169
|
+
print(s)
|
|
1170
|
+
|
|
1171
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
|
|
1172
|
+
|
|
1173
|
+
Prints:
|
|
1174
|
+
|
|
1175
|
+
.. code-block:: text
|
|
1176
|
+
|
|
1177
|
+
[64] = tile(shape=(1), storage=register)
|
|
1178
|
+
|
|
1179
|
+
|
|
1180
|
+
"""
|
|
1181
|
+
...
|
|
1182
|
+
|
|
1183
|
+
@over
|
|
1184
|
+
def tile_argmin(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Int, Tuple[1]]:
|
|
1185
|
+
"""Cooperatively compute the index of the minimum element in the tile using all threads in the block.
|
|
1186
|
+
|
|
1187
|
+
:param a: The tile to compute the argmin from
|
|
1188
|
+
:returns: A single-element tile holding the index of the minimum value
|
|
1189
|
+
|
|
1190
|
+
Example:
|
|
1191
|
+
|
|
1192
|
+
.. code-block:: python
|
|
1193
|
+
|
|
1194
|
+
@wp.kernel
|
|
1195
|
+
def compute():
|
|
1196
|
+
t = wp.tile_arange(64, 128)
|
|
1197
|
+
s = wp.tile_argmin(t)
|
|
1198
|
+
|
|
1199
|
+
print(s)
|
|
1200
|
+
|
|
1201
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
|
|
1202
|
+
|
|
1203
|
+
Prints:
|
|
1204
|
+
|
|
1205
|
+
.. code-block:: text
|
|
1206
|
+
|
|
1207
|
+
[0] = tile(shape=(1), storage=register)
|
|
1208
|
+
|
|
1209
|
+
|
|
1210
|
+
"""
|
|
1211
|
+
...
|
|
1212
|
+
|
|
1213
|
+
@over
|
|
1214
|
+
def tile_max(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[1]]:
|
|
1215
|
+
"""Cooperatively compute the maximum of the tile elements using all threads in the block.
|
|
1216
|
+
|
|
1217
|
+
:param a: The tile to compute the maximum from
|
|
1218
|
+
:returns: A single-element tile holding the maximum value
|
|
1219
|
+
|
|
1220
|
+
Example:
|
|
1221
|
+
|
|
1222
|
+
.. code-block:: python
|
|
1223
|
+
|
|
1224
|
+
@wp.kernel
|
|
1225
|
+
def compute():
|
|
1226
|
+
t = wp.tile_arange(64, 128)
|
|
1227
|
+
s = wp.tile_max(t)
|
|
1228
|
+
|
|
1229
|
+
print(s)
|
|
1230
|
+
|
|
1231
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
|
|
1232
|
+
|
|
1233
|
+
Prints:
|
|
1234
|
+
|
|
1235
|
+
.. code-block:: text
|
|
1236
|
+
|
|
1237
|
+
[127] = tile(shape=(1), storage=register)
|
|
1238
|
+
|
|
1239
|
+
|
|
1240
|
+
"""
|
|
1241
|
+
...
|
|
1242
|
+
|
|
1243
|
+
@over
|
|
1244
|
+
def tile_argmax(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Int, Tuple[1]]:
|
|
1245
|
+
"""Cooperatively compute the index of the maximum element in the tile using all threads in the block.
|
|
1246
|
+
|
|
1247
|
+
:param a: The tile to compute the argmax from
|
|
1248
|
+
:returns: A single-element tile holding the index of the maximum value
|
|
1249
|
+
|
|
1250
|
+
Example:
|
|
1251
|
+
|
|
1252
|
+
.. code-block:: python
|
|
1253
|
+
|
|
1254
|
+
@wp.kernel
|
|
1255
|
+
def compute():
|
|
1256
|
+
t = wp.tile_arange(64, 128)
|
|
1257
|
+
s = wp.tile_argmax(t)
|
|
1258
|
+
|
|
1259
|
+
print(s)
|
|
1260
|
+
|
|
1261
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=64)
|
|
1262
|
+
|
|
1263
|
+
Prints:
|
|
1264
|
+
|
|
1265
|
+
.. code-block:: text
|
|
1266
|
+
|
|
1267
|
+
[63] = tile(shape=(1), storage=register)
|
|
1268
|
+
|
|
1269
|
+
|
|
1270
|
+
"""
|
|
1271
|
+
...
|
|
1272
|
+
|
|
1273
|
+
@over
|
|
1274
|
+
def tile_reduce(op: Callable, a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[1]]:
|
|
1275
|
+
"""Apply a custom reduction operator across the tile.
|
|
1276
|
+
|
|
1277
|
+
This function cooperatively performs a reduction using the provided operator across the tile.
|
|
1278
|
+
|
|
1279
|
+
:param op: A callable function that accepts two arguments and returns one argument, may be a user function or builtin
|
|
1280
|
+
:param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's data type
|
|
1281
|
+
:returns: A single-element tile with the same data type as the input tile.
|
|
1282
|
+
|
|
1283
|
+
Example:
|
|
1284
|
+
|
|
1285
|
+
.. code-block:: python
|
|
1286
|
+
|
|
1287
|
+
@wp.kernel
|
|
1288
|
+
def factorial():
|
|
1289
|
+
t = wp.tile_arange(1, 10, dtype=int)
|
|
1290
|
+
s = wp.tile_reduce(wp.mul, t)
|
|
1291
|
+
|
|
1292
|
+
print(s)
|
|
1293
|
+
|
|
1294
|
+
wp.launch_tiled(factorial, dim=[1], inputs=[], block_dim=16)
|
|
1295
|
+
|
|
1296
|
+
Prints:
|
|
1297
|
+
|
|
1298
|
+
.. code-block:: text
|
|
1299
|
+
|
|
1300
|
+
[362880] = tile(shape=(1), storage=register)
|
|
1301
|
+
|
|
1302
|
+
"""
|
|
1303
|
+
...
|
|
1304
|
+
|
|
1305
|
+
@over
|
|
1306
|
+
def tile_scan_inclusive(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
|
|
1307
|
+
"""Inclusive scan (prefix sum) across the tile.
|
|
1308
|
+
|
|
1309
|
+
This function cooperatively performs an inclusive scan (cumulative sum) across the tile.
|
|
1310
|
+
|
|
1311
|
+
:param a: The input tile. Must be a tile of type float32, int32, or uint32.
|
|
1312
|
+
:returns: A new tile containing the inclusive scan result.
|
|
1313
|
+
|
|
1314
|
+
Example:
|
|
1315
|
+
|
|
1316
|
+
.. code-block:: python
|
|
1317
|
+
|
|
1318
|
+
@wp.kernel
|
|
1319
|
+
def scan_example():
|
|
1320
|
+
t = wp.tile_arange(1, 5, dtype=int)
|
|
1321
|
+
s = wp.tile_scan_inclusive(t)
|
|
1322
|
+
print(s)
|
|
1323
|
+
|
|
1324
|
+
wp.launch_tiled(scan_example, dim=[1], inputs=[], block_dim=16)
|
|
1325
|
+
|
|
1326
|
+
Prints:
|
|
1327
|
+
|
|
1328
|
+
.. code-block:: text
|
|
1329
|
+
|
|
1330
|
+
[1, 3, 6, 10] = tile(shape=(4), storage=register)
|
|
1331
|
+
|
|
1332
|
+
"""
|
|
1333
|
+
...
|
|
1334
|
+
|
|
1335
|
+
@over
|
|
1336
|
+
def tile_scan_exclusive(a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
|
|
1337
|
+
"""Exclusive scan (prefix sum) across the tile.
|
|
1338
|
+
|
|
1339
|
+
This function cooperatively performs an exclusive scan (cumulative sum) across the tile.
|
|
1340
|
+
|
|
1341
|
+
:param a: The input tile. Must be a tile of type float32, int32, or uint32.
|
|
1342
|
+
:returns: A new tile containing the exclusive scan result.
|
|
1343
|
+
|
|
1344
|
+
Example:
|
|
1345
|
+
|
|
1346
|
+
.. code-block:: python
|
|
1347
|
+
|
|
1348
|
+
@wp.kernel
|
|
1349
|
+
def scan_example():
|
|
1350
|
+
t = wp.tile_arange(1, 5, dtype=int)
|
|
1351
|
+
s = wp.tile_scan_exclusive(t)
|
|
1352
|
+
print(s)
|
|
1353
|
+
|
|
1354
|
+
wp.launch_tiled(scan_example, dim=[1], inputs=[], block_dim=16)
|
|
1355
|
+
|
|
1356
|
+
Prints:
|
|
1357
|
+
|
|
1358
|
+
.. code-block:: text
|
|
1359
|
+
|
|
1360
|
+
[0, 1, 3, 6] = tile(shape=(4), storage=register)
|
|
1361
|
+
|
|
1362
|
+
"""
|
|
1363
|
+
...
|
|
1364
|
+
|
|
1365
|
+
@over
|
|
1366
|
+
def tile_map(op: Callable, a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
|
|
1367
|
+
"""Apply a unary function onto the tile.
|
|
1368
|
+
|
|
1369
|
+
This function cooperatively applies a unary function to each element of the tile using all threads in the block.
|
|
1370
|
+
|
|
1371
|
+
:param op: A callable function that accepts one argument and returns one argument, may be a user function or builtin
|
|
1372
|
+
:param a: The input tile, the operator (or one of its overloads) must be able to accept the tile's data type
|
|
1373
|
+
:returns: A tile with the same dimensions and data type as the input tile.
|
|
1374
|
+
|
|
1375
|
+
Example:
|
|
1376
|
+
|
|
1377
|
+
.. code-block:: python
|
|
1378
|
+
|
|
1379
|
+
@wp.kernel
|
|
1380
|
+
def compute():
|
|
1381
|
+
t = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
|
|
1382
|
+
s = wp.tile_map(wp.sin, t)
|
|
1383
|
+
|
|
1384
|
+
print(s)
|
|
1385
|
+
|
|
1386
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=16)
|
|
1387
|
+
|
|
1388
|
+
Prints:
|
|
1389
|
+
|
|
1390
|
+
.. code-block:: text
|
|
1391
|
+
|
|
1392
|
+
[0 0.0998334 0.198669 0.29552 0.389418 0.479426 0.564642 0.644218 0.717356 0.783327] = tile(shape=(10), storage=register)
|
|
1393
|
+
|
|
1394
|
+
"""
|
|
1395
|
+
...
|
|
1396
|
+
|
|
1397
|
+
@over
|
|
1398
|
+
def tile_map(
|
|
1399
|
+
op: Callable, a: Tile[Scalar, Tuple[int, ...]], b: Tile[Scalar, Tuple[int, ...]]
|
|
1400
|
+
) -> Tile[Scalar, Tuple[int, ...]]:
|
|
1401
|
+
"""Apply a binary function onto the tile.
|
|
1402
|
+
|
|
1403
|
+
This function cooperatively applies a binary function to each element of the tiles using all threads in the block.
|
|
1404
|
+
Both input tiles must have the same dimensions and datatype.
|
|
1405
|
+
|
|
1406
|
+
:param op: A callable function that accepts two arguments and returns one argument, all of the same type, may be a user function or builtin
|
|
1407
|
+
:param a: The first input tile, the operator (or one of its overloads) must be able to accept the tile's dtype
|
|
1408
|
+
:param b: The second input tile, the operator (or one of its overloads) must be able to accept the tile's dtype
|
|
1409
|
+
:returns: A tile with the same dimensions and datatype as the input tiles.
|
|
1410
|
+
|
|
1411
|
+
Example:
|
|
1412
|
+
|
|
1413
|
+
.. code-block:: python
|
|
1414
|
+
|
|
1415
|
+
@wp.kernel
|
|
1416
|
+
def compute():
|
|
1417
|
+
a = wp.tile_arange(0.0, 1.0, 0.1, dtype=float)
|
|
1418
|
+
b = wp.tile_ones(shape=10, dtype=float)
|
|
1419
|
+
|
|
1420
|
+
s = wp.tile_map(wp.add, a, b)
|
|
1421
|
+
|
|
1422
|
+
print(s)
|
|
1423
|
+
|
|
1424
|
+
wp.launch_tiled(compute, dim=[1], inputs=[], block_dim=16)
|
|
1425
|
+
|
|
1426
|
+
Prints:
|
|
1427
|
+
|
|
1428
|
+
.. code-block:: text
|
|
1429
|
+
|
|
1430
|
+
[1 1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 1.9] = tile(shape=(10), storage=register)
|
|
1431
|
+
"""
|
|
1432
|
+
...
|
|
1433
|
+
|
|
1434
|
+
@over
|
|
1435
|
+
def bvh_query_aabb(id: uint64, low: vec3f, high: vec3f) -> BvhQuery:
|
|
1436
|
+
"""Construct an axis-aligned bounding box query against a BVH object.
|
|
1437
|
+
|
|
1438
|
+
This query can be used to iterate over all bounds inside a BVH.
|
|
1439
|
+
|
|
1440
|
+
:param id: The BVH identifier
|
|
1441
|
+
:param low: The lower bound of the bounding box in BVH space
|
|
1442
|
+
:param high: The upper bound of the bounding box in BVH space
|
|
1443
|
+
"""
|
|
1444
|
+
...
|
|
1445
|
+
|
|
1446
|
+
@over
|
|
1447
|
+
def bvh_query_ray(id: uint64, start: vec3f, dir: vec3f) -> BvhQuery:
|
|
1448
|
+
"""Construct a ray query against a BVH object.
|
|
1449
|
+
|
|
1450
|
+
This query can be used to iterate over all bounds that intersect the ray.
|
|
1451
|
+
|
|
1452
|
+
:param id: The BVH identifier
|
|
1453
|
+
:param start: The start of the ray in BVH space
|
|
1454
|
+
:param dir: The direction of the ray in BVH space
|
|
1455
|
+
"""
|
|
1456
|
+
...
|
|
1457
|
+
|
|
1458
|
+
@over
|
|
1459
|
+
def bvh_query_next(query: BvhQuery, index: int32) -> bool:
|
|
1460
|
+
"""Move to the next bound returned by the query.
|
|
1461
|
+
The index of the current bound is stored in ``index``, returns ``False`` if there are no more overlapping bound.
|
|
1462
|
+
"""
|
|
1463
|
+
...
|
|
1464
|
+
|
|
1465
|
+
@over
|
|
1466
|
+
def mesh_query_point(id: uint64, point: vec3f, max_dist: float32) -> MeshQueryPoint:
|
|
1467
|
+
"""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
|
|
1468
|
+
|
|
1469
|
+
Identifies the sign of the distance using additional ray-casts to determine if the point is inside or outside.
|
|
1470
|
+
This method is relatively robust, but does increase computational cost.
|
|
1471
|
+
See below for additional sign determination methods.
|
|
1472
|
+
|
|
1473
|
+
:param id: The mesh identifier
|
|
1474
|
+
:param point: The point in space to query
|
|
1475
|
+
:param max_dist: Mesh faces above this distance will not be considered by the query
|
|
1476
|
+
"""
|
|
1477
|
+
...
|
|
1478
|
+
|
|
1479
|
+
@over
|
|
1480
|
+
def mesh_query_point_no_sign(id: uint64, point: vec3f, max_dist: float32) -> MeshQueryPoint:
|
|
1481
|
+
"""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
|
|
1482
|
+
|
|
1483
|
+
This method does not compute the sign of the point (inside/outside) which makes it faster than other point query methods.
|
|
1484
|
+
|
|
1485
|
+
:param id: The mesh identifier
|
|
1486
|
+
:param point: The point in space to query
|
|
1487
|
+
:param max_dist: Mesh faces above this distance will not be considered by the query
|
|
1488
|
+
"""
|
|
1489
|
+
...
|
|
1490
|
+
|
|
1491
|
+
@over
|
|
1492
|
+
def mesh_query_furthest_point_no_sign(id: uint64, point: vec3f, min_dist: float32) -> MeshQueryPoint:
|
|
1493
|
+
"""Computes the furthest point on the mesh with identifier `id` to the given point in space.
|
|
1494
|
+
|
|
1495
|
+
This method does not compute the sign of the point (inside/outside).
|
|
1496
|
+
|
|
1497
|
+
:param id: The mesh identifier
|
|
1498
|
+
:param point: The point in space to query
|
|
1499
|
+
:param min_dist: Mesh faces below this distance will not be considered by the query
|
|
1500
|
+
"""
|
|
1501
|
+
...
|
|
1502
|
+
|
|
1503
|
+
@over
|
|
1504
|
+
def mesh_query_point_sign_normal(id: uint64, point: vec3f, max_dist: float32, epsilon: float32) -> MeshQueryPoint:
|
|
1505
|
+
"""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given ``point`` in space.
|
|
1506
|
+
|
|
1507
|
+
Identifies the sign of the distance (inside/outside) using the angle-weighted pseudo normal.
|
|
1508
|
+
This approach to sign determination is robust for well conditioned meshes that are watertight and non-self intersecting.
|
|
1509
|
+
It is also comparatively fast to compute.
|
|
1510
|
+
|
|
1511
|
+
:param id: The mesh identifier
|
|
1512
|
+
:param point: The point in space to query
|
|
1513
|
+
:param max_dist: Mesh faces above this distance will not be considered by the query
|
|
1514
|
+
:param epsilon: Epsilon treating distance values as equal, when locating the minimum distance vertex/face/edge, as a
|
|
1515
|
+
fraction of the average edge length, also for treating closest point as being on edge/vertex default 1e-3
|
|
1516
|
+
"""
|
|
1517
|
+
...
|
|
1518
|
+
|
|
1519
|
+
@over
|
|
1520
|
+
def mesh_query_point_sign_winding_number(
|
|
1521
|
+
id: uint64, point: vec3f, max_dist: float32, accuracy: float32, threshold: float32
|
|
1522
|
+
) -> MeshQueryPoint:
|
|
1523
|
+
"""Computes the closest point on the :class:`Mesh` with identifier ``id`` to the given point in space.
|
|
1524
|
+
|
|
1525
|
+
Identifies the sign using the winding number of the mesh relative to the query point. This method of sign determination is robust for poorly conditioned meshes
|
|
1526
|
+
and provides a smooth approximation to sign even when the mesh is not watertight. This method is the most robust and accurate of the sign determination meshes
|
|
1527
|
+
but also the most expensive.
|
|
1528
|
+
|
|
1529
|
+
.. note:: The :class:`Mesh` object must be constructed with ``support_winding_number=True`` for this method to return correct results.
|
|
1530
|
+
|
|
1531
|
+
:param id: The mesh identifier
|
|
1532
|
+
:param point: The point in space to query
|
|
1533
|
+
:param max_dist: Mesh faces above this distance will not be considered by the query
|
|
1534
|
+
:param accuracy: Accuracy for computing the winding number with fast winding number method utilizing second-order dipole approximation, default 2.0
|
|
1535
|
+
:param threshold: The threshold of the winding number to be considered inside, default 0.5
|
|
1536
|
+
"""
|
|
1537
|
+
...
|
|
1538
|
+
|
|
1539
|
+
@over
|
|
1540
|
+
def mesh_query_ray(id: uint64, start: vec3f, dir: vec3f, max_t: float32) -> MeshQueryRay:
|
|
1541
|
+
"""Computes the closest ray hit on the :class:`Mesh` with identifier ``id``.
|
|
1542
|
+
|
|
1543
|
+
:param id: The mesh identifier
|
|
1544
|
+
:param start: The start point of the ray
|
|
1545
|
+
:param dir: The ray direction (should be normalized)
|
|
1546
|
+
:param max_t: The maximum distance along the ray to check for intersections
|
|
1547
|
+
"""
|
|
1548
|
+
...
|
|
1549
|
+
|
|
1550
|
+
@over
|
|
1551
|
+
def mesh_query_aabb(id: uint64, low: vec3f, high: vec3f) -> MeshQueryAABB:
|
|
1552
|
+
"""Construct an axis-aligned bounding box query against a :class:`Mesh`.
|
|
1553
|
+
|
|
1554
|
+
This query can be used to iterate over all triangles inside a volume.
|
|
1555
|
+
|
|
1556
|
+
:param id: The mesh identifier
|
|
1557
|
+
:param low: The lower bound of the bounding box in mesh space
|
|
1558
|
+
:param high: The upper bound of the bounding box in mesh space
|
|
1559
|
+
"""
|
|
1560
|
+
...
|
|
1561
|
+
|
|
1562
|
+
@over
|
|
1563
|
+
def mesh_query_aabb_next(query: MeshQueryAABB, index: int32) -> bool:
|
|
1564
|
+
"""Move to the next triangle overlapping the query bounding box.
|
|
1565
|
+
|
|
1566
|
+
The index of the current face is stored in ``index``, returns ``False`` if there are no more overlapping triangles.
|
|
1567
|
+
"""
|
|
1568
|
+
...
|
|
1569
|
+
|
|
1570
|
+
@over
|
|
1571
|
+
def mesh_eval_position(id: uint64, face: int32, bary_u: float32, bary_v: float32) -> vec3f:
|
|
1572
|
+
"""Evaluates the position on the :class:`Mesh` given a face index and barycentric coordinates."""
|
|
1573
|
+
...
|
|
1574
|
+
|
|
1575
|
+
@over
|
|
1576
|
+
def mesh_eval_velocity(id: uint64, face: int32, bary_u: float32, bary_v: float32) -> vec3f:
|
|
1577
|
+
"""Evaluates the velocity on the :class:`Mesh` given a face index and barycentric coordinates."""
|
|
1578
|
+
...
|
|
1579
|
+
|
|
1580
|
+
@over
|
|
1581
|
+
def hash_grid_query(id: uint64, point: vec3f, max_dist: float32) -> HashGridQuery:
|
|
1582
|
+
"""Construct a point query against a :class:`HashGrid`.
|
|
1583
|
+
|
|
1584
|
+
This query can be used to iterate over all neighboring point within a fixed radius from the query point.
|
|
1585
|
+
"""
|
|
1586
|
+
...
|
|
1587
|
+
|
|
1588
|
+
@over
|
|
1589
|
+
def hash_grid_query_next(query: HashGridQuery, index: int32) -> bool:
|
|
1590
|
+
"""Move to the next point in the hash grid query.
|
|
1591
|
+
|
|
1592
|
+
The index of the current neighbor is stored in ``index``, returns ``False`` if there are no more neighbors.
|
|
1593
|
+
"""
|
|
1594
|
+
...
|
|
1595
|
+
|
|
1596
|
+
@over
|
|
1597
|
+
def hash_grid_point_id(id: uint64, index: int32) -> int:
|
|
1598
|
+
"""Return the index of a point in the :class:`HashGrid`.
|
|
1599
|
+
|
|
1600
|
+
This can be used to reorder threads such that grid traversal occurs in a spatially coherent order.
|
|
1601
|
+
|
|
1602
|
+
Returns -1 if the :class:`HashGrid` has not been reserved.
|
|
1603
|
+
"""
|
|
1604
|
+
...
|
|
1605
|
+
|
|
1606
|
+
@over
|
|
1607
|
+
def intersect_tri_tri(v0: vec3f, v1: vec3f, v2: vec3f, u0: vec3f, u1: vec3f, u2: vec3f) -> int:
|
|
1608
|
+
"""Tests for intersection between two triangles (v0, v1, v2) and (u0, u1, u2) using Moller's method.
|
|
1609
|
+
|
|
1610
|
+
Returns > 0 if triangles intersect.
|
|
1611
|
+
"""
|
|
1612
|
+
...
|
|
1613
|
+
|
|
1614
|
+
@over
|
|
1615
|
+
def mesh_get(id: uint64) -> Mesh:
|
|
1616
|
+
"""Retrieves the mesh given its index."""
|
|
1617
|
+
...
|
|
1618
|
+
|
|
1619
|
+
@over
|
|
1620
|
+
def mesh_eval_face_normal(id: uint64, face: int32) -> vec3f:
|
|
1621
|
+
"""Evaluates the face normal the mesh given a face index."""
|
|
1622
|
+
...
|
|
1623
|
+
|
|
1624
|
+
@over
|
|
1625
|
+
def mesh_get_point(id: uint64, index: int32) -> vec3f:
|
|
1626
|
+
"""Returns the point of the mesh given a index."""
|
|
1627
|
+
...
|
|
1628
|
+
|
|
1629
|
+
@over
|
|
1630
|
+
def mesh_get_velocity(id: uint64, index: int32) -> vec3f:
|
|
1631
|
+
"""Returns the velocity of the mesh given a index."""
|
|
1632
|
+
...
|
|
1633
|
+
|
|
1634
|
+
@over
|
|
1635
|
+
def mesh_get_index(id: uint64, index: int32) -> int:
|
|
1636
|
+
"""Returns the point-index of the mesh given a face-vertex index."""
|
|
1637
|
+
...
|
|
1638
|
+
|
|
1639
|
+
@over
|
|
1640
|
+
def closest_point_edge_edge(p1: vec3f, q1: vec3f, p2: vec3f, q2: vec3f, epsilon: float32) -> vec3f:
|
|
1641
|
+
"""Finds the closest points between two edges.
|
|
1642
|
+
|
|
1643
|
+
Returns barycentric weights to the points on each edge, as well as the closest distance between the edges.
|
|
1644
|
+
|
|
1645
|
+
:param p1: First point of first edge
|
|
1646
|
+
:param q1: Second point of first edge
|
|
1647
|
+
:param p2: First point of second edge
|
|
1648
|
+
:param q2: Second point of second edge
|
|
1649
|
+
:param epsilon: Zero tolerance for determining if points in an edge are degenerate.
|
|
1650
|
+
:param out: vec3 output containing (s,t,d), where `s` in [0,1] is the barycentric weight for the first edge, `t` is the barycentric weight for the second edge, and `d` is the distance between the two edges at these two closest points.
|
|
1651
|
+
"""
|
|
1652
|
+
...
|
|
1653
|
+
|
|
1654
|
+
@over
|
|
1655
|
+
def reversed(range: range_t) -> range_t:
|
|
1656
|
+
"""Returns the range in reversed order."""
|
|
1657
|
+
...
|
|
1658
|
+
|
|
1659
|
+
@over
|
|
1660
|
+
def volume_sample(id: uint64, uvw: vec3f, sampling_mode: int32, dtype: Any) -> Any:
|
|
1661
|
+
"""Sample the volume of type `dtype` given by ``id`` at the volume local-space point ``uvw``.
|
|
1662
|
+
|
|
1663
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
|
|
1664
|
+
"""
|
|
1665
|
+
...
|
|
1666
|
+
|
|
1667
|
+
@over
|
|
1668
|
+
def volume_sample_grad(id: uint64, uvw: vec3f, sampling_mode: int32, grad: Any, dtype: Any) -> Any:
|
|
1669
|
+
"""Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
|
|
1670
|
+
|
|
1671
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
|
|
1672
|
+
"""
|
|
1673
|
+
...
|
|
1674
|
+
|
|
1675
|
+
@over
|
|
1676
|
+
def volume_lookup(id: uint64, i: int32, j: int32, k: int32, dtype: Any) -> Any:
|
|
1677
|
+
"""Returns the value of voxel with coordinates ``i``, ``j``, ``k`` for a volume of type type `dtype`.
|
|
1678
|
+
|
|
1679
|
+
If the voxel at this index does not exist, this function returns the background value.
|
|
1680
|
+
"""
|
|
1681
|
+
...
|
|
1682
|
+
|
|
1683
|
+
@over
|
|
1684
|
+
def volume_store(id: uint64, i: int32, j: int32, k: int32, value: Any):
|
|
1685
|
+
"""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``."""
|
|
1686
|
+
...
|
|
1687
|
+
|
|
1688
|
+
@over
|
|
1689
|
+
def volume_sample_f(id: uint64, uvw: vec3f, sampling_mode: int32) -> float:
|
|
1690
|
+
"""Sample the volume given by ``id`` at the volume local-space point ``uvw``.
|
|
1691
|
+
|
|
1692
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
|
|
1693
|
+
"""
|
|
1694
|
+
...
|
|
1695
|
+
|
|
1696
|
+
@over
|
|
1697
|
+
def volume_sample_grad_f(id: uint64, uvw: vec3f, sampling_mode: int32, grad: vec3f) -> float:
|
|
1698
|
+
"""Sample the volume and its gradient given by ``id`` at the volume local-space point ``uvw``.
|
|
1699
|
+
|
|
1700
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
|
|
1701
|
+
"""
|
|
1702
|
+
...
|
|
1703
|
+
|
|
1704
|
+
@over
|
|
1705
|
+
def volume_lookup_f(id: uint64, i: int32, j: int32, k: int32) -> float:
|
|
1706
|
+
"""Returns the value of voxel with coordinates ``i``, ``j``, ``k``.
|
|
1707
|
+
|
|
1708
|
+
If the voxel at this index does not exist, this function returns the background value
|
|
1709
|
+
"""
|
|
1710
|
+
...
|
|
1711
|
+
|
|
1712
|
+
@over
|
|
1713
|
+
def volume_store_f(id: uint64, i: int32, j: int32, k: int32, value: float32):
|
|
1714
|
+
"""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``."""
|
|
1715
|
+
...
|
|
1716
|
+
|
|
1717
|
+
@over
|
|
1718
|
+
def volume_sample_v(id: uint64, uvw: vec3f, sampling_mode: int32) -> vec3f:
|
|
1719
|
+
"""Sample the vector volume given by ``id`` at the volume local-space point ``uvw``.
|
|
1720
|
+
|
|
1721
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
|
|
1722
|
+
"""
|
|
1723
|
+
...
|
|
1724
|
+
|
|
1725
|
+
@over
|
|
1726
|
+
def volume_lookup_v(id: uint64, i: int32, j: int32, k: int32) -> vec3f:
|
|
1727
|
+
"""Returns the vector value of voxel with coordinates ``i``, ``j``, ``k``.
|
|
1728
|
+
|
|
1729
|
+
If the voxel at this index does not exist, this function returns the background value.
|
|
1730
|
+
"""
|
|
1731
|
+
...
|
|
1732
|
+
|
|
1733
|
+
@over
|
|
1734
|
+
def volume_store_v(id: uint64, i: int32, j: int32, k: int32, value: vec3f):
|
|
1735
|
+
"""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``."""
|
|
1736
|
+
...
|
|
1737
|
+
|
|
1738
|
+
@over
|
|
1739
|
+
def volume_sample_i(id: uint64, uvw: vec3f) -> int:
|
|
1740
|
+
"""Sample the :class:`int32` volume given by ``id`` at the volume local-space point ``uvw``."""
|
|
1741
|
+
...
|
|
1742
|
+
|
|
1743
|
+
@over
|
|
1744
|
+
def volume_lookup_i(id: uint64, i: int32, j: int32, k: int32) -> int:
|
|
1745
|
+
"""Returns the :class:`int32` value of voxel with coordinates ``i``, ``j``, ``k``.
|
|
1746
|
+
|
|
1747
|
+
If the voxel at this index does not exist, this function returns the background value.
|
|
1748
|
+
"""
|
|
1749
|
+
...
|
|
1750
|
+
|
|
1751
|
+
@over
|
|
1752
|
+
def volume_store_i(id: uint64, i: int32, j: int32, k: int32, value: int32):
|
|
1753
|
+
"""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``."""
|
|
1754
|
+
...
|
|
1755
|
+
|
|
1756
|
+
@over
|
|
1757
|
+
def volume_sample_index(id: uint64, uvw: vec3f, sampling_mode: int32, voxel_data: Array[Any], background: Any) -> Any:
|
|
1758
|
+
"""Sample the volume given by ``id`` at the volume local-space point ``uvw``.
|
|
1759
|
+
|
|
1760
|
+
Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
|
|
1761
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
|
|
1762
|
+
This function is available for both index grids and classical volumes.
|
|
1763
|
+
|
|
1764
|
+
"""
|
|
1765
|
+
...
|
|
1766
|
+
|
|
1767
|
+
@over
|
|
1768
|
+
def volume_sample_grad_index(
|
|
1769
|
+
id: uint64, uvw: vec3f, sampling_mode: int32, voxel_data: Array[Any], background: Any, grad: Any
|
|
1770
|
+
) -> Any:
|
|
1771
|
+
"""Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
|
|
1772
|
+
|
|
1773
|
+
Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
|
|
1774
|
+
Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
|
|
1775
|
+
This function is available for both index grids and classical volumes.
|
|
1776
|
+
|
|
1777
|
+
"""
|
|
1778
|
+
...
|
|
1779
|
+
|
|
1780
|
+
@over
|
|
1781
|
+
def volume_lookup_index(id: uint64, i: int32, j: int32, k: int32) -> int32:
|
|
1782
|
+
"""Returns the index associated to the voxel with coordinates ``i``, ``j``, ``k``.
|
|
1783
|
+
|
|
1784
|
+
If the voxel at this index does not exist, this function returns -1.
|
|
1785
|
+
This function is available for both index grids and classical volumes.
|
|
1786
|
+
|
|
1787
|
+
"""
|
|
1788
|
+
...
|
|
1789
|
+
|
|
1790
|
+
@over
|
|
1791
|
+
def volume_index_to_world(id: uint64, uvw: vec3f) -> vec3f:
|
|
1792
|
+
"""Transform a point ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation."""
|
|
1793
|
+
...
|
|
1794
|
+
|
|
1795
|
+
@over
|
|
1796
|
+
def volume_world_to_index(id: uint64, xyz: vec3f) -> vec3f:
|
|
1797
|
+
"""Transform a point ``xyz`` defined in volume world space to the volume's index space given the volume's intrinsic affine transformation."""
|
|
1798
|
+
...
|
|
1799
|
+
|
|
1800
|
+
@over
|
|
1801
|
+
def volume_index_to_world_dir(id: uint64, uvw: vec3f) -> vec3f:
|
|
1802
|
+
"""Transform a direction ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation."""
|
|
1803
|
+
...
|
|
1804
|
+
|
|
1805
|
+
@over
|
|
1806
|
+
def volume_world_to_index_dir(id: uint64, xyz: vec3f) -> vec3f:
|
|
1807
|
+
"""Transform a direction ``xyz`` defined in volume world space to the volume's index space given the volume's intrinsic affine transformation."""
|
|
1808
|
+
...
|
|
1809
|
+
|
|
1810
|
+
@over
|
|
1811
|
+
def rand_init(seed: int32) -> uint32:
|
|
1812
|
+
"""Initialize a new random number generator given a user-defined seed. Returns a 32-bit integer representing the RNG state."""
|
|
1813
|
+
...
|
|
1814
|
+
|
|
1815
|
+
@over
|
|
1816
|
+
def rand_init(seed: int32, offset: int32) -> uint32:
|
|
1817
|
+
"""Initialize a new random number generator given a user-defined seed and an offset.
|
|
1818
|
+
|
|
1819
|
+
This alternative constructor can be useful in parallel programs, where a kernel as a whole should share a seed,
|
|
1820
|
+
but each thread should generate uncorrelated values. In this case usage should be ``r = rand_init(seed, tid)``
|
|
1821
|
+
"""
|
|
1822
|
+
...
|
|
1823
|
+
|
|
1824
|
+
@over
|
|
1825
|
+
def randi(state: uint32) -> int:
|
|
1826
|
+
"""Return a random integer in the range [-2^31, 2^31)."""
|
|
1827
|
+
...
|
|
1828
|
+
|
|
1829
|
+
@over
|
|
1830
|
+
def randi(state: uint32, low: int32, high: int32) -> int:
|
|
1831
|
+
"""Return a random integer between [low, high)."""
|
|
1832
|
+
...
|
|
1833
|
+
|
|
1834
|
+
@over
|
|
1835
|
+
def randu(state: uint32) -> uint32:
|
|
1836
|
+
"""Return a random unsigned integer in the range [0, 2^32)."""
|
|
1837
|
+
...
|
|
1838
|
+
|
|
1839
|
+
@over
|
|
1840
|
+
def randu(state: uint32, low: uint32, high: uint32) -> uint32:
|
|
1841
|
+
"""Return a random unsigned integer between [low, high)."""
|
|
1842
|
+
...
|
|
1843
|
+
|
|
1844
|
+
@over
|
|
1845
|
+
def randf(state: uint32) -> float:
|
|
1846
|
+
"""Return a random float between [0.0, 1.0)."""
|
|
1847
|
+
...
|
|
1848
|
+
|
|
1849
|
+
@over
|
|
1850
|
+
def randf(state: uint32, low: float32, high: float32) -> float:
|
|
1851
|
+
"""Return a random float between [low, high)."""
|
|
1852
|
+
...
|
|
1853
|
+
|
|
1854
|
+
@over
|
|
1855
|
+
def randn(state: uint32) -> float:
|
|
1856
|
+
"""Sample a normal (Gaussian) distribution of mean 0 and variance 1."""
|
|
1857
|
+
...
|
|
1858
|
+
|
|
1859
|
+
@over
|
|
1860
|
+
def sample_cdf(state: uint32, cdf: Array[float32]) -> int:
|
|
1861
|
+
"""Inverse-transform sample a cumulative distribution function."""
|
|
1862
|
+
...
|
|
1863
|
+
|
|
1864
|
+
@over
|
|
1865
|
+
def sample_triangle(state: uint32) -> vec2f:
|
|
1866
|
+
"""Uniformly sample a triangle. Returns sample barycentric coordinates."""
|
|
1867
|
+
...
|
|
1868
|
+
|
|
1869
|
+
@over
|
|
1870
|
+
def sample_unit_ring(state: uint32) -> vec2f:
|
|
1871
|
+
"""Uniformly sample a ring in the xy plane."""
|
|
1872
|
+
...
|
|
1873
|
+
|
|
1874
|
+
@over
|
|
1875
|
+
def sample_unit_disk(state: uint32) -> vec2f:
|
|
1876
|
+
"""Uniformly sample a disk in the xy plane."""
|
|
1877
|
+
...
|
|
1878
|
+
|
|
1879
|
+
@over
|
|
1880
|
+
def sample_unit_sphere_surface(state: uint32) -> vec3f:
|
|
1881
|
+
"""Uniformly sample a unit sphere surface."""
|
|
1882
|
+
...
|
|
1883
|
+
|
|
1884
|
+
@over
|
|
1885
|
+
def sample_unit_sphere(state: uint32) -> vec3f:
|
|
1886
|
+
"""Uniformly sample a unit sphere."""
|
|
1887
|
+
...
|
|
1888
|
+
|
|
1889
|
+
@over
|
|
1890
|
+
def sample_unit_hemisphere_surface(state: uint32) -> vec3f:
|
|
1891
|
+
"""Uniformly sample a unit hemisphere surface."""
|
|
1892
|
+
...
|
|
1893
|
+
|
|
1894
|
+
@over
|
|
1895
|
+
def sample_unit_hemisphere(state: uint32) -> vec3f:
|
|
1896
|
+
"""Uniformly sample a unit hemisphere."""
|
|
1897
|
+
...
|
|
1898
|
+
|
|
1899
|
+
@over
|
|
1900
|
+
def sample_unit_square(state: uint32) -> vec2f:
|
|
1901
|
+
"""Uniformly sample a unit square."""
|
|
1902
|
+
...
|
|
1903
|
+
|
|
1904
|
+
@over
|
|
1905
|
+
def sample_unit_cube(state: uint32) -> vec3f:
|
|
1906
|
+
"""Uniformly sample a unit cube."""
|
|
1907
|
+
...
|
|
1908
|
+
|
|
1909
|
+
@over
|
|
1910
|
+
def poisson(state: uint32, lam: float32) -> uint32:
|
|
1911
|
+
"""Generate a random sample from a Poisson distribution.
|
|
1912
|
+
|
|
1913
|
+
:param state: RNG state
|
|
1914
|
+
:param lam: The expected value of the distribution
|
|
1915
|
+
"""
|
|
1916
|
+
...
|
|
1917
|
+
|
|
1918
|
+
@over
|
|
1919
|
+
def noise(state: uint32, x: float32) -> float:
|
|
1920
|
+
"""Non-periodic Perlin-style noise in 1D."""
|
|
1921
|
+
...
|
|
1922
|
+
|
|
1923
|
+
@over
|
|
1924
|
+
def noise(state: uint32, xy: vec2f) -> float:
|
|
1925
|
+
"""Non-periodic Perlin-style noise in 2D."""
|
|
1926
|
+
...
|
|
1927
|
+
|
|
1928
|
+
@over
|
|
1929
|
+
def noise(state: uint32, xyz: vec3f) -> float:
|
|
1930
|
+
"""Non-periodic Perlin-style noise in 3D."""
|
|
1931
|
+
...
|
|
1932
|
+
|
|
1933
|
+
@over
|
|
1934
|
+
def noise(state: uint32, xyzt: vec4f) -> float:
|
|
1935
|
+
"""Non-periodic Perlin-style noise in 4D."""
|
|
1936
|
+
...
|
|
1937
|
+
|
|
1938
|
+
@over
|
|
1939
|
+
def pnoise(state: uint32, x: float32, px: int32) -> float:
|
|
1940
|
+
"""Periodic Perlin-style noise in 1D."""
|
|
1941
|
+
...
|
|
1942
|
+
|
|
1943
|
+
@over
|
|
1944
|
+
def pnoise(state: uint32, xy: vec2f, px: int32, py: int32) -> float:
|
|
1945
|
+
"""Periodic Perlin-style noise in 2D."""
|
|
1946
|
+
...
|
|
1947
|
+
|
|
1948
|
+
@over
|
|
1949
|
+
def pnoise(state: uint32, xyz: vec3f, px: int32, py: int32, pz: int32) -> float:
|
|
1950
|
+
"""Periodic Perlin-style noise in 3D."""
|
|
1951
|
+
...
|
|
1952
|
+
|
|
1953
|
+
@over
|
|
1954
|
+
def pnoise(state: uint32, xyzt: vec4f, px: int32, py: int32, pz: int32, pt: int32) -> float:
|
|
1955
|
+
"""Periodic Perlin-style noise in 4D."""
|
|
1956
|
+
...
|
|
1957
|
+
|
|
1958
|
+
@over
|
|
1959
|
+
def curlnoise(state: uint32, xy: vec2f, octaves: uint32, lacunarity: float32, gain: float32) -> vec2f:
|
|
1960
|
+
"""Divergence-free vector field based on the gradient of a Perlin noise function."""
|
|
1961
|
+
...
|
|
1962
|
+
|
|
1963
|
+
@over
|
|
1964
|
+
def curlnoise(state: uint32, xyz: vec3f, octaves: uint32, lacunarity: float32, gain: float32) -> vec3f:
|
|
1965
|
+
"""Divergence-free vector field based on the curl of three Perlin noise functions."""
|
|
1966
|
+
...
|
|
1967
|
+
|
|
1968
|
+
@over
|
|
1969
|
+
def curlnoise(state: uint32, xyzt: vec4f, octaves: uint32, lacunarity: float32, gain: float32) -> vec3f:
|
|
1970
|
+
"""Divergence-free vector field based on the curl of three Perlin noise functions."""
|
|
1971
|
+
...
|
|
1972
|
+
|
|
1973
|
+
@over
|
|
1974
|
+
def printf(fmt: str, *args: Any):
|
|
1975
|
+
"""Allows printing formatted strings using C-style format specifiers."""
|
|
1976
|
+
...
|
|
1977
|
+
|
|
1978
|
+
@over
|
|
1979
|
+
def print(value: Any):
|
|
1980
|
+
"""Print variable to stdout"""
|
|
1981
|
+
...
|
|
1982
|
+
|
|
1983
|
+
@over
|
|
1984
|
+
def breakpoint():
|
|
1985
|
+
"""Debugger breakpoint"""
|
|
1986
|
+
...
|
|
1987
|
+
|
|
1988
|
+
@over
|
|
1989
|
+
def tid() -> int:
|
|
1990
|
+
"""Return the current thread index for a 1D kernel launch.
|
|
1991
|
+
|
|
1992
|
+
Note that this is the *global* index of the thread in the range [0, dim)
|
|
1993
|
+
where dim is the parameter passed to kernel launch.
|
|
1994
|
+
|
|
1995
|
+
This function may not be called from user-defined Warp functions.
|
|
1996
|
+
"""
|
|
1997
|
+
...
|
|
1998
|
+
|
|
1999
|
+
@over
|
|
2000
|
+
def tid() -> Tuple[int, int]:
|
|
2001
|
+
"""Return the current thread indices for a 2D kernel launch.
|
|
2002
|
+
|
|
2003
|
+
Use ``i,j = wp.tid()`` syntax to retrieve the coordinates inside the kernel thread grid.
|
|
2004
|
+
|
|
2005
|
+
This function may not be called from user-defined Warp functions.
|
|
2006
|
+
"""
|
|
2007
|
+
...
|
|
2008
|
+
|
|
2009
|
+
@over
|
|
2010
|
+
def tid() -> Tuple[int, int, int]:
|
|
2011
|
+
"""Return the current thread indices for a 3D kernel launch.
|
|
2012
|
+
|
|
2013
|
+
Use ``i,j,k = wp.tid()`` syntax to retrieve the coordinates inside the kernel thread grid.
|
|
2014
|
+
|
|
2015
|
+
This function may not be called from user-defined Warp functions.
|
|
2016
|
+
"""
|
|
2017
|
+
...
|
|
2018
|
+
|
|
2019
|
+
@over
|
|
2020
|
+
def tid() -> Tuple[int, int, int, int]:
|
|
2021
|
+
"""Return the current thread indices for a 4D kernel launch.
|
|
2022
|
+
|
|
2023
|
+
Use ``i,j,k,l = wp.tid()`` syntax to retrieve the coordinates inside the kernel thread grid.
|
|
2024
|
+
|
|
2025
|
+
This function may not be called from user-defined Warp functions.
|
|
2026
|
+
"""
|
|
2027
|
+
...
|
|
2028
|
+
|
|
2029
|
+
@over
|
|
2030
|
+
def block_dim() -> int:
|
|
2031
|
+
"""Returns the number of threads in the current block."""
|
|
2032
|
+
...
|
|
2033
|
+
|
|
2034
|
+
@over
|
|
2035
|
+
def select(cond: bool, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2036
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2037
|
+
|
|
2038
|
+
.. deprecated:: 1.7
|
|
2039
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2040
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2041
|
+
"""
|
|
2042
|
+
...
|
|
2043
|
+
|
|
2044
|
+
@over
|
|
2045
|
+
def select(cond: int8, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2046
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2047
|
+
|
|
2048
|
+
.. deprecated:: 1.7
|
|
2049
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2050
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2051
|
+
"""
|
|
2052
|
+
...
|
|
2053
|
+
|
|
2054
|
+
@over
|
|
2055
|
+
def select(cond: uint8, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2056
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2057
|
+
|
|
2058
|
+
.. deprecated:: 1.7
|
|
2059
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2060
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2061
|
+
"""
|
|
2062
|
+
...
|
|
2063
|
+
|
|
2064
|
+
@over
|
|
2065
|
+
def select(cond: int16, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2066
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2067
|
+
|
|
2068
|
+
.. deprecated:: 1.7
|
|
2069
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2070
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2071
|
+
"""
|
|
2072
|
+
...
|
|
2073
|
+
|
|
2074
|
+
@over
|
|
2075
|
+
def select(cond: uint16, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2076
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2077
|
+
|
|
2078
|
+
.. deprecated:: 1.7
|
|
2079
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2080
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2081
|
+
"""
|
|
2082
|
+
...
|
|
2083
|
+
|
|
2084
|
+
@over
|
|
2085
|
+
def select(cond: int32, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2086
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2087
|
+
|
|
2088
|
+
.. deprecated:: 1.7
|
|
2089
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2090
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2091
|
+
"""
|
|
2092
|
+
...
|
|
2093
|
+
|
|
2094
|
+
@over
|
|
2095
|
+
def select(cond: uint32, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2096
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2097
|
+
|
|
2098
|
+
.. deprecated:: 1.7
|
|
2099
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2100
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2101
|
+
"""
|
|
2102
|
+
...
|
|
2103
|
+
|
|
2104
|
+
@over
|
|
2105
|
+
def select(cond: int64, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2106
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2107
|
+
|
|
2108
|
+
.. deprecated:: 1.7
|
|
2109
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2110
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2111
|
+
"""
|
|
2112
|
+
...
|
|
2113
|
+
|
|
2114
|
+
@over
|
|
2115
|
+
def select(cond: uint64, value_if_false: Any, value_if_true: Any) -> Any:
|
|
2116
|
+
"""Select between two arguments, if ``cond`` is ``False`` then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2117
|
+
|
|
2118
|
+
.. deprecated:: 1.7
|
|
2119
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2120
|
+
``where(cond, value_if_true, value_if_false)``.
|
|
2121
|
+
"""
|
|
2122
|
+
...
|
|
2123
|
+
|
|
2124
|
+
@over
|
|
2125
|
+
def select(arr: Array[Any], value_if_false: Any, value_if_true: Any) -> Any:
|
|
2126
|
+
"""Select between two arguments, if ``arr`` is null then return ``value_if_false``, otherwise return ``value_if_true``.
|
|
2127
|
+
|
|
2128
|
+
.. deprecated:: 1.7
|
|
2129
|
+
Use :func:`where` instead, which has the more intuitive argument order:
|
|
2130
|
+
``where(arr, value_if_true, value_if_false)``.
|
|
2131
|
+
"""
|
|
2132
|
+
...
|
|
2133
|
+
|
|
2134
|
+
@over
|
|
2135
|
+
def where(cond: bool, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2136
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2137
|
+
...
|
|
2138
|
+
|
|
2139
|
+
@over
|
|
2140
|
+
def where(cond: int8, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2141
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2142
|
+
...
|
|
2143
|
+
|
|
2144
|
+
@over
|
|
2145
|
+
def where(cond: uint8, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2146
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2147
|
+
...
|
|
2148
|
+
|
|
2149
|
+
@over
|
|
2150
|
+
def where(cond: int16, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2151
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2152
|
+
...
|
|
2153
|
+
|
|
2154
|
+
@over
|
|
2155
|
+
def where(cond: uint16, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2156
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2157
|
+
...
|
|
2158
|
+
|
|
2159
|
+
@over
|
|
2160
|
+
def where(cond: int32, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2161
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2162
|
+
...
|
|
2163
|
+
|
|
2164
|
+
@over
|
|
2165
|
+
def where(cond: uint32, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2166
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2167
|
+
...
|
|
2168
|
+
|
|
2169
|
+
@over
|
|
2170
|
+
def where(cond: int64, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2171
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2172
|
+
...
|
|
2173
|
+
|
|
2174
|
+
@over
|
|
2175
|
+
def where(cond: uint64, value_if_true: Any, value_if_false: Any) -> Any:
|
|
2176
|
+
"""Select between two arguments, if ``cond`` is ``True`` then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2177
|
+
...
|
|
2178
|
+
|
|
2179
|
+
@over
|
|
2180
|
+
def where(arr: Array[Any], value_if_true: Any, value_if_false: Any) -> Any:
|
|
2181
|
+
"""Select between two arguments, if ``arr`` is not null then return ``value_if_true``, otherwise return ``value_if_false``."""
|
|
2182
|
+
...
|
|
2183
|
+
|
|
2184
|
+
@over
|
|
2185
|
+
def atomic_add(arr: Array[Any], i: Int, value: Any) -> Any:
|
|
2186
|
+
"""Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
|
|
2187
|
+
This function is automatically invoked when using the syntax ``arr[i] += value``.
|
|
2188
|
+
"""
|
|
2189
|
+
...
|
|
2190
|
+
|
|
2191
|
+
@over
|
|
2192
|
+
def atomic_add(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2193
|
+
"""Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
|
|
2194
|
+
This function is automatically invoked when using the syntax ``arr[i,j] += value``.
|
|
2195
|
+
"""
|
|
2196
|
+
...
|
|
2197
|
+
|
|
2198
|
+
@over
|
|
2199
|
+
def atomic_add(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2200
|
+
"""Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
|
|
2201
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.
|
|
2202
|
+
"""
|
|
2203
|
+
...
|
|
2204
|
+
|
|
2205
|
+
@over
|
|
2206
|
+
def atomic_add(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2207
|
+
"""Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
|
|
2208
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.
|
|
2209
|
+
"""
|
|
2210
|
+
...
|
|
2211
|
+
|
|
2212
|
+
@over
|
|
2213
|
+
def atomic_add(arr: FabricArray[Any], i: Int, value: Any) -> Any:
|
|
2214
|
+
"""Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
|
|
2215
|
+
This function is automatically invoked when using the syntax ``arr[i] += value``.
|
|
2216
|
+
"""
|
|
2217
|
+
...
|
|
2218
|
+
|
|
2219
|
+
@over
|
|
2220
|
+
def atomic_add(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2221
|
+
"""Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
|
|
2222
|
+
This function is automatically invoked when using the syntax ``arr[i,j] += value``.
|
|
2223
|
+
"""
|
|
2224
|
+
...
|
|
2225
|
+
|
|
2226
|
+
@over
|
|
2227
|
+
def atomic_add(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2228
|
+
"""Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
|
|
2229
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.
|
|
2230
|
+
"""
|
|
2231
|
+
...
|
|
2232
|
+
|
|
2233
|
+
@over
|
|
2234
|
+
def atomic_add(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2235
|
+
"""Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
|
|
2236
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.
|
|
2237
|
+
"""
|
|
2238
|
+
...
|
|
2239
|
+
|
|
2240
|
+
@over
|
|
2241
|
+
def atomic_add(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
|
|
2242
|
+
"""Atomically adds ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
|
|
2243
|
+
This function is automatically invoked when using the syntax ``arr[i] += value``.
|
|
2244
|
+
"""
|
|
2245
|
+
...
|
|
2246
|
+
|
|
2247
|
+
@over
|
|
2248
|
+
def atomic_add(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2249
|
+
"""Atomically adds ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
|
|
2250
|
+
This function is automatically invoked when using the syntax ``arr[i,j] += value``.
|
|
2251
|
+
"""
|
|
2252
|
+
...
|
|
2253
|
+
|
|
2254
|
+
@over
|
|
2255
|
+
def atomic_add(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2256
|
+
"""Atomically adds ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
|
|
2257
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k] += value``.
|
|
2258
|
+
"""
|
|
2259
|
+
...
|
|
2260
|
+
|
|
2261
|
+
@over
|
|
2262
|
+
def atomic_add(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2263
|
+
"""Atomically adds ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
|
|
2264
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k,l] += value``.
|
|
2265
|
+
"""
|
|
2266
|
+
...
|
|
2267
|
+
|
|
2268
|
+
@over
|
|
2269
|
+
def atomic_sub(arr: Array[Any], i: Int, value: Any) -> Any:
|
|
2270
|
+
"""Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
|
|
2271
|
+
This function is automatically invoked when using the syntax ``arr[i] -= value``.
|
|
2272
|
+
"""
|
|
2273
|
+
...
|
|
2274
|
+
|
|
2275
|
+
@over
|
|
2276
|
+
def atomic_sub(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2277
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
|
|
2278
|
+
This function is automatically invoked when using the syntax ``arr[i,j] -= value``.
|
|
2279
|
+
"""
|
|
2280
|
+
...
|
|
2281
|
+
|
|
2282
|
+
@over
|
|
2283
|
+
def atomic_sub(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2284
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
|
|
2285
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.
|
|
2286
|
+
"""
|
|
2287
|
+
...
|
|
2288
|
+
|
|
2289
|
+
@over
|
|
2290
|
+
def atomic_sub(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2291
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
|
|
2292
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.
|
|
2293
|
+
"""
|
|
2294
|
+
...
|
|
2295
|
+
|
|
2296
|
+
@over
|
|
2297
|
+
def atomic_sub(arr: FabricArray[Any], i: Int, value: Any) -> Any:
|
|
2298
|
+
"""Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
|
|
2299
|
+
This function is automatically invoked when using the syntax ``arr[i] -= value``.
|
|
2300
|
+
"""
|
|
2301
|
+
...
|
|
2302
|
+
|
|
2303
|
+
@over
|
|
2304
|
+
def atomic_sub(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2305
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
|
|
2306
|
+
This function is automatically invoked when using the syntax ``arr[i,j] -= value``.
|
|
2307
|
+
"""
|
|
2308
|
+
...
|
|
2309
|
+
|
|
2310
|
+
@over
|
|
2311
|
+
def atomic_sub(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2312
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
|
|
2313
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.
|
|
2314
|
+
"""
|
|
2315
|
+
...
|
|
2316
|
+
|
|
2317
|
+
@over
|
|
2318
|
+
def atomic_sub(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2319
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
|
|
2320
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.
|
|
2321
|
+
"""
|
|
2322
|
+
...
|
|
2323
|
+
|
|
2324
|
+
@over
|
|
2325
|
+
def atomic_sub(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
|
|
2326
|
+
"""Atomically subtracts ``value`` onto ``arr[i]`` and returns the original value of ``arr[i]``.
|
|
2327
|
+
This function is automatically invoked when using the syntax ``arr[i] -= value``.
|
|
2328
|
+
"""
|
|
2329
|
+
...
|
|
2330
|
+
|
|
2331
|
+
@over
|
|
2332
|
+
def atomic_sub(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2333
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j]`` and returns the original value of ``arr[i,j]``.
|
|
2334
|
+
This function is automatically invoked when using the syntax ``arr[i,j] -= value``.
|
|
2335
|
+
"""
|
|
2336
|
+
...
|
|
2337
|
+
|
|
2338
|
+
@over
|
|
2339
|
+
def atomic_sub(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2340
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j,k]`` and returns the original value of ``arr[i,j,k]``.
|
|
2341
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k] -= value``.
|
|
2342
|
+
"""
|
|
2343
|
+
...
|
|
2344
|
+
|
|
2345
|
+
@over
|
|
2346
|
+
def atomic_sub(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2347
|
+
"""Atomically subtracts ``value`` onto ``arr[i,j,k,l]`` and returns the original value of ``arr[i,j,k,l]``.
|
|
2348
|
+
This function is automatically invoked when using the syntax ``arr[i,j,k,l] -= value``.
|
|
2349
|
+
"""
|
|
2350
|
+
...
|
|
2351
|
+
|
|
2352
|
+
@over
|
|
2353
|
+
def atomic_min(arr: Array[Any], i: Int, value: Any) -> Any:
|
|
2354
|
+
"""Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
|
|
2355
|
+
|
|
2356
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2357
|
+
"""
|
|
2358
|
+
...
|
|
2359
|
+
|
|
2360
|
+
@over
|
|
2361
|
+
def atomic_min(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2362
|
+
"""Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
|
|
2363
|
+
|
|
2364
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2365
|
+
"""
|
|
2366
|
+
...
|
|
2367
|
+
|
|
2368
|
+
@over
|
|
2369
|
+
def atomic_min(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2370
|
+
"""Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
|
|
2371
|
+
|
|
2372
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2373
|
+
"""
|
|
2374
|
+
...
|
|
2375
|
+
|
|
2376
|
+
@over
|
|
2377
|
+
def atomic_min(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2378
|
+
"""Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
|
|
2379
|
+
|
|
2380
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2381
|
+
"""
|
|
2382
|
+
...
|
|
2383
|
+
|
|
2384
|
+
@over
|
|
2385
|
+
def atomic_min(arr: FabricArray[Any], i: Int, value: Any) -> Any:
|
|
2386
|
+
"""Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
|
|
2387
|
+
|
|
2388
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2389
|
+
"""
|
|
2390
|
+
...
|
|
2391
|
+
|
|
2392
|
+
@over
|
|
2393
|
+
def atomic_min(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2394
|
+
"""Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
|
|
2395
|
+
|
|
2396
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2397
|
+
"""
|
|
2398
|
+
...
|
|
2399
|
+
|
|
2400
|
+
@over
|
|
2401
|
+
def atomic_min(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2402
|
+
"""Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
|
|
2403
|
+
|
|
2404
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2405
|
+
"""
|
|
2406
|
+
...
|
|
2407
|
+
|
|
2408
|
+
@over
|
|
2409
|
+
def atomic_min(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2410
|
+
"""Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
|
|
2411
|
+
|
|
2412
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2413
|
+
"""
|
|
2414
|
+
...
|
|
2415
|
+
|
|
2416
|
+
@over
|
|
2417
|
+
def atomic_min(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
|
|
2418
|
+
"""Compute the minimum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
|
|
2419
|
+
|
|
2420
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2421
|
+
"""
|
|
2422
|
+
...
|
|
2423
|
+
|
|
2424
|
+
@over
|
|
2425
|
+
def atomic_min(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2426
|
+
"""Compute the minimum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
|
|
2427
|
+
|
|
2428
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2429
|
+
"""
|
|
2430
|
+
...
|
|
2431
|
+
|
|
2432
|
+
@over
|
|
2433
|
+
def atomic_min(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2434
|
+
"""Compute the minimum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
|
|
2435
|
+
|
|
2436
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2437
|
+
"""
|
|
2438
|
+
...
|
|
2439
|
+
|
|
2440
|
+
@over
|
|
2441
|
+
def atomic_min(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2442
|
+
"""Compute the minimum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
|
|
2443
|
+
|
|
2444
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2445
|
+
"""
|
|
2446
|
+
...
|
|
2447
|
+
|
|
2448
|
+
@over
|
|
2449
|
+
def atomic_max(arr: Array[Any], i: Int, value: Any) -> Any:
|
|
2450
|
+
"""Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
|
|
2451
|
+
|
|
2452
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2453
|
+
"""
|
|
2454
|
+
...
|
|
2455
|
+
|
|
2456
|
+
@over
|
|
2457
|
+
def atomic_max(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2458
|
+
"""Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
|
|
2459
|
+
|
|
2460
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2461
|
+
"""
|
|
2462
|
+
...
|
|
2463
|
+
|
|
2464
|
+
@over
|
|
2465
|
+
def atomic_max(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2466
|
+
"""Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
|
|
2467
|
+
|
|
2468
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2469
|
+
"""
|
|
2470
|
+
...
|
|
2471
|
+
|
|
2472
|
+
@over
|
|
2473
|
+
def atomic_max(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2474
|
+
"""Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
|
|
2475
|
+
|
|
2476
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2477
|
+
"""
|
|
2478
|
+
...
|
|
2479
|
+
|
|
2480
|
+
@over
|
|
2481
|
+
def atomic_max(arr: FabricArray[Any], i: Int, value: Any) -> Any:
|
|
2482
|
+
"""Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
|
|
2483
|
+
|
|
2484
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2485
|
+
"""
|
|
2486
|
+
...
|
|
2487
|
+
|
|
2488
|
+
@over
|
|
2489
|
+
def atomic_max(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2490
|
+
"""Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
|
|
2491
|
+
|
|
2492
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2493
|
+
"""
|
|
2494
|
+
...
|
|
2495
|
+
|
|
2496
|
+
@over
|
|
2497
|
+
def atomic_max(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2498
|
+
"""Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
|
|
2499
|
+
|
|
2500
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2501
|
+
"""
|
|
2502
|
+
...
|
|
2503
|
+
|
|
2504
|
+
@over
|
|
2505
|
+
def atomic_max(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2506
|
+
"""Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
|
|
2507
|
+
|
|
2508
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2509
|
+
"""
|
|
2510
|
+
...
|
|
2511
|
+
|
|
2512
|
+
@over
|
|
2513
|
+
def atomic_max(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
|
|
2514
|
+
"""Compute the maximum of ``value`` and ``arr[i]``, atomically update the array, and return the old value.
|
|
2515
|
+
|
|
2516
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2517
|
+
"""
|
|
2518
|
+
...
|
|
2519
|
+
|
|
2520
|
+
@over
|
|
2521
|
+
def atomic_max(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2522
|
+
"""Compute the maximum of ``value`` and ``arr[i,j]``, atomically update the array, and return the old value.
|
|
2523
|
+
|
|
2524
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2525
|
+
"""
|
|
2526
|
+
...
|
|
2527
|
+
|
|
2528
|
+
@over
|
|
2529
|
+
def atomic_max(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2530
|
+
"""Compute the maximum of ``value`` and ``arr[i,j,k]``, atomically update the array, and return the old value.
|
|
2531
|
+
|
|
2532
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2533
|
+
"""
|
|
2534
|
+
...
|
|
2535
|
+
|
|
2536
|
+
@over
|
|
2537
|
+
def atomic_max(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2538
|
+
"""Compute the maximum of ``value`` and ``arr[i,j,k,l]``, atomically update the array, and return the old value.
|
|
2539
|
+
|
|
2540
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2541
|
+
"""
|
|
2542
|
+
...
|
|
2543
|
+
|
|
2544
|
+
@over
|
|
2545
|
+
def atomic_cas(arr: Array[Any], i: Int, compare: Any, value: Any) -> Any:
|
|
2546
|
+
"""Atomically compare and swap ``value`` with ``arr[i]`` if ``arr[i]`` equals ``compare``, and return the old value.
|
|
2547
|
+
|
|
2548
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2549
|
+
"""
|
|
2550
|
+
...
|
|
2551
|
+
|
|
2552
|
+
@over
|
|
2553
|
+
def atomic_cas(arr: Array[Any], i: Int, j: Int, compare: Any, value: Any) -> Any:
|
|
2554
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j]`` if ``arr[i,j]`` equals ``compare``, and return the old value.
|
|
2555
|
+
|
|
2556
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2557
|
+
"""
|
|
2558
|
+
...
|
|
2559
|
+
|
|
2560
|
+
@over
|
|
2561
|
+
def atomic_cas(arr: Array[Any], i: Int, j: Int, k: Int, compare: Any, value: Any) -> Any:
|
|
2562
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j,k]`` if ``arr[i,j,k]`` equals ``compare``, and return the old value.
|
|
2563
|
+
|
|
2564
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2565
|
+
"""
|
|
2566
|
+
...
|
|
2567
|
+
|
|
2568
|
+
@over
|
|
2569
|
+
def atomic_cas(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, compare: Any, value: Any) -> Any:
|
|
2570
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j,k,l]`` if ``arr[i,j,k,l]`` equals ``compare``, and return the old value.
|
|
2571
|
+
|
|
2572
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2573
|
+
"""
|
|
2574
|
+
...
|
|
2575
|
+
|
|
2576
|
+
@over
|
|
2577
|
+
def atomic_cas(arr: FabricArray[Any], i: Int, compare: Any, value: Any) -> Any:
|
|
2578
|
+
"""Atomically compare and swap ``value`` with ``arr[i]`` if ``arr[i]`` equals ``compare``, and return the old value.
|
|
2579
|
+
|
|
2580
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2581
|
+
"""
|
|
2582
|
+
...
|
|
2583
|
+
|
|
2584
|
+
@over
|
|
2585
|
+
def atomic_cas(arr: FabricArray[Any], i: Int, j: Int, compare: Any, value: Any) -> Any:
|
|
2586
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j]`` if ``arr[i,j]`` equals ``compare``, and return the old value.
|
|
2587
|
+
|
|
2588
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2589
|
+
"""
|
|
2590
|
+
...
|
|
2591
|
+
|
|
2592
|
+
@over
|
|
2593
|
+
def atomic_cas(arr: FabricArray[Any], i: Int, j: Int, k: Int, compare: Any, value: Any) -> Any:
|
|
2594
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j,k]`` if ``arr[i,j,k]`` equals ``compare``, and return the old value.
|
|
2595
|
+
|
|
2596
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2597
|
+
"""
|
|
2598
|
+
...
|
|
2599
|
+
|
|
2600
|
+
@over
|
|
2601
|
+
def atomic_cas(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, compare: Any, value: Any) -> Any:
|
|
2602
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j,k,l]`` if ``arr[i,j,k,l]`` equals ``compare``, and return the old value.
|
|
2603
|
+
|
|
2604
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2605
|
+
"""
|
|
2606
|
+
...
|
|
2607
|
+
|
|
2608
|
+
@over
|
|
2609
|
+
def atomic_cas(arr: IndexedFabricArray[Any], i: Int, compare: Any, value: Any) -> Any:
|
|
2610
|
+
"""Atomically compare and swap ``value`` with ``arr[i]`` if ``arr[i]`` equals ``compare``, and return the old value.
|
|
2611
|
+
|
|
2612
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2613
|
+
"""
|
|
2614
|
+
...
|
|
2615
|
+
|
|
2616
|
+
@over
|
|
2617
|
+
def atomic_cas(arr: IndexedFabricArray[Any], i: Int, j: Int, compare: Any, value: Any) -> Any:
|
|
2618
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j]`` if ``arr[i,j]`` equals ``compare``, and return the old value.
|
|
2619
|
+
|
|
2620
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2621
|
+
"""
|
|
2622
|
+
...
|
|
2623
|
+
|
|
2624
|
+
@over
|
|
2625
|
+
def atomic_cas(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, compare: Any, value: Any) -> Any:
|
|
2626
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j,k]`` if ``arr[i,j,k]`` equals ``compare``, and return the old value.
|
|
2627
|
+
|
|
2628
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2629
|
+
"""
|
|
2630
|
+
...
|
|
2631
|
+
|
|
2632
|
+
@over
|
|
2633
|
+
def atomic_cas(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, compare: Any, value: Any) -> Any:
|
|
2634
|
+
"""Atomically compare and swap ``value`` with ``arr[i,j,k,l]`` if ``arr[i,j,k,l]`` equals ``compare``, and return the old value.
|
|
2635
|
+
|
|
2636
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2637
|
+
"""
|
|
2638
|
+
...
|
|
2639
|
+
|
|
2640
|
+
@over
|
|
2641
|
+
def atomic_exch(arr: Array[Any], i: Int, value: Any) -> Any:
|
|
2642
|
+
"""Atomically exchange ``value`` with ``arr[i]`` and return the old value.
|
|
2643
|
+
|
|
2644
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2645
|
+
"""
|
|
2646
|
+
...
|
|
2647
|
+
|
|
2648
|
+
@over
|
|
2649
|
+
def atomic_exch(arr: Array[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2650
|
+
"""Atomically exchange ``value`` with ``arr[i,j]`` and return the old value.
|
|
2651
|
+
|
|
2652
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2653
|
+
"""
|
|
2654
|
+
...
|
|
2655
|
+
|
|
2656
|
+
@over
|
|
2657
|
+
def atomic_exch(arr: Array[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2658
|
+
"""Atomically exchange ``value`` with ``arr[i,j,k]`` and return the old value.
|
|
2659
|
+
|
|
2660
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2661
|
+
"""
|
|
2662
|
+
...
|
|
2663
|
+
|
|
2664
|
+
@over
|
|
2665
|
+
def atomic_exch(arr: Array[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2666
|
+
"""Atomically exchange ``value`` with ``arr[i,j,k,l]`` and return the old value.
|
|
2667
|
+
|
|
2668
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2669
|
+
"""
|
|
2670
|
+
...
|
|
2671
|
+
|
|
2672
|
+
@over
|
|
2673
|
+
def atomic_exch(arr: FabricArray[Any], i: Int, value: Any) -> Any:
|
|
2674
|
+
"""Atomically exchange ``value`` with ``arr[i]`` and return the old value.
|
|
2675
|
+
|
|
2676
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2677
|
+
"""
|
|
2678
|
+
...
|
|
2679
|
+
|
|
2680
|
+
@over
|
|
2681
|
+
def atomic_exch(arr: FabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2682
|
+
"""Atomically exchange ``value`` with ``arr[i,j]`` and return the old value.
|
|
2683
|
+
|
|
2684
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2685
|
+
"""
|
|
2686
|
+
...
|
|
2687
|
+
|
|
2688
|
+
@over
|
|
2689
|
+
def atomic_exch(arr: FabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2690
|
+
"""Atomically exchange ``value`` with ``arr[i,j,k]`` and return the old value.
|
|
2691
|
+
|
|
2692
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2693
|
+
"""
|
|
2694
|
+
...
|
|
2695
|
+
|
|
2696
|
+
@over
|
|
2697
|
+
def atomic_exch(arr: FabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2698
|
+
"""Atomically exchange ``value`` with ``arr[i,j,k,l]`` and return the old value.
|
|
2699
|
+
|
|
2700
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2701
|
+
"""
|
|
2702
|
+
...
|
|
2703
|
+
|
|
2704
|
+
@over
|
|
2705
|
+
def atomic_exch(arr: IndexedFabricArray[Any], i: Int, value: Any) -> Any:
|
|
2706
|
+
"""Atomically exchange ``value`` with ``arr[i]`` and return the old value.
|
|
2707
|
+
|
|
2708
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2709
|
+
"""
|
|
2710
|
+
...
|
|
2711
|
+
|
|
2712
|
+
@over
|
|
2713
|
+
def atomic_exch(arr: IndexedFabricArray[Any], i: Int, j: Int, value: Any) -> Any:
|
|
2714
|
+
"""Atomically exchange ``value`` with ``arr[i,j]`` and return the old value.
|
|
2715
|
+
|
|
2716
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2717
|
+
"""
|
|
2718
|
+
...
|
|
2719
|
+
|
|
2720
|
+
@over
|
|
2721
|
+
def atomic_exch(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, value: Any) -> Any:
|
|
2722
|
+
"""Atomically exchange ``value`` with ``arr[i,j,k]`` and return the old value.
|
|
2723
|
+
|
|
2724
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2725
|
+
"""
|
|
2726
|
+
...
|
|
2727
|
+
|
|
2728
|
+
@over
|
|
2729
|
+
def atomic_exch(arr: IndexedFabricArray[Any], i: Int, j: Int, k: Int, l: Int, value: Any) -> Any:
|
|
2730
|
+
"""Atomically exchange ``value`` with ``arr[i,j,k,l]`` and return the old value.
|
|
2731
|
+
|
|
2732
|
+
The operation is only atomic on a per-component basis for vectors and matrices.
|
|
2733
|
+
"""
|
|
2734
|
+
...
|
|
2735
|
+
|
|
2736
|
+
@over
|
|
2737
|
+
def lerp(a: Float, b: Float, t: Float) -> Float:
|
|
2738
|
+
"""Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``"""
|
|
2739
|
+
...
|
|
2740
|
+
|
|
2741
|
+
@over
|
|
2742
|
+
def lerp(a: Vector[Any, Float], b: Vector[Any, Float], t: Float) -> Vector[Any, Float]:
|
|
2743
|
+
"""Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``"""
|
|
2744
|
+
...
|
|
2745
|
+
|
|
2746
|
+
@over
|
|
2747
|
+
def lerp(a: Matrix[Any, Any, Float], b: Matrix[Any, Any, Float], t: Float) -> Matrix[Any, Any, Float]:
|
|
2748
|
+
"""Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``"""
|
|
2749
|
+
...
|
|
2750
|
+
|
|
2751
|
+
@over
|
|
2752
|
+
def lerp(a: Quaternion[Float], b: Quaternion[Float], t: Float) -> Quaternion[Float]:
|
|
2753
|
+
"""Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``"""
|
|
2754
|
+
...
|
|
2755
|
+
|
|
2756
|
+
@over
|
|
2757
|
+
def lerp(a: Transformation[Float], b: Transformation[Float], t: Float) -> Transformation[Float]:
|
|
2758
|
+
"""Linearly interpolate two values ``a`` and ``b`` using factor ``t``, computed as ``a*(1-t) + b*t``"""
|
|
2759
|
+
...
|
|
2760
|
+
|
|
2761
|
+
@over
|
|
2762
|
+
def smoothstep(a: Float, b: Float, x: Float) -> Float:
|
|
2763
|
+
"""Smoothly interpolate between two values ``a`` and ``b`` using a factor ``x``,
|
|
2764
|
+
and return a result between 0 and 1 using a cubic Hermite interpolation after clamping.
|
|
2765
|
+
"""
|
|
2766
|
+
...
|
|
2767
|
+
|
|
2768
|
+
@over
|
|
2769
|
+
def expect_near(a: Float, b: Float, tolerance: Float):
|
|
2770
|
+
"""Prints an error to stdout if ``a`` and ``b`` are not closer than tolerance in magnitude"""
|
|
2771
|
+
...
|
|
2772
|
+
|
|
2773
|
+
@over
|
|
2774
|
+
def expect_near(a: Vector[Any, Float], b: Vector[Any, Float], tolerance: Float):
|
|
2775
|
+
"""Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude"""
|
|
2776
|
+
...
|
|
2777
|
+
|
|
2778
|
+
@over
|
|
2779
|
+
def expect_near(a: Quaternion[Float], b: Quaternion[Float], tolerance: Float):
|
|
2780
|
+
"""Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude"""
|
|
2781
|
+
...
|
|
2782
|
+
|
|
2783
|
+
@over
|
|
2784
|
+
def expect_near(a: Matrix[Any, Any, Float], b: Matrix[Any, Any, Float], tolerance: Float):
|
|
2785
|
+
"""Prints an error to stdout if any element of ``a`` and ``b`` are not closer than tolerance in magnitude"""
|
|
2786
|
+
...
|
|
2787
|
+
|
|
2788
|
+
@over
|
|
2789
|
+
def lower_bound(arr: Array[Scalar], value: Scalar) -> int:
|
|
2790
|
+
"""Search a sorted array ``arr`` for the closest element greater than or equal to ``value``."""
|
|
2791
|
+
...
|
|
2792
|
+
|
|
2793
|
+
@over
|
|
2794
|
+
def lower_bound(arr: Array[Scalar], arr_begin: int32, arr_end: int32, value: Scalar) -> int:
|
|
2795
|
+
"""Search a sorted array ``arr`` in the range [arr_begin, arr_end) for the closest element greater than or equal to ``value``."""
|
|
2796
|
+
...
|
|
2797
|
+
|
|
2798
|
+
@over
|
|
2799
|
+
def add(a: Scalar, b: Scalar) -> Scalar:
|
|
2800
|
+
""" """
|
|
2801
|
+
...
|
|
2802
|
+
|
|
2803
|
+
@over
|
|
2804
|
+
def add(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2805
|
+
""" """
|
|
2806
|
+
...
|
|
2807
|
+
|
|
2808
|
+
@over
|
|
2809
|
+
def add(a: Quaternion[Scalar], b: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
2810
|
+
""" """
|
|
2811
|
+
...
|
|
2812
|
+
|
|
2813
|
+
@over
|
|
2814
|
+
def add(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
2815
|
+
""" """
|
|
2816
|
+
...
|
|
2817
|
+
|
|
2818
|
+
@over
|
|
2819
|
+
def add(a: Transformation[Scalar], b: Transformation[Scalar]) -> Transformation[Scalar]:
|
|
2820
|
+
""" """
|
|
2821
|
+
...
|
|
2822
|
+
|
|
2823
|
+
@over
|
|
2824
|
+
def add(a: Tile[Any, Tuple[int, ...]], b: Tile[Any, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
|
|
2825
|
+
"""Add each element of two tiles together"""
|
|
2826
|
+
...
|
|
2827
|
+
|
|
2828
|
+
@over
|
|
2829
|
+
def sub(a: Scalar, b: Scalar) -> Scalar:
|
|
2830
|
+
""" """
|
|
2831
|
+
...
|
|
2832
|
+
|
|
2833
|
+
@over
|
|
2834
|
+
def sub(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2835
|
+
""" """
|
|
2836
|
+
...
|
|
2837
|
+
|
|
2838
|
+
@over
|
|
2839
|
+
def sub(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
2840
|
+
""" """
|
|
2841
|
+
...
|
|
2842
|
+
|
|
2843
|
+
@over
|
|
2844
|
+
def sub(a: Quaternion[Scalar], b: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
2845
|
+
""" """
|
|
2846
|
+
...
|
|
2847
|
+
|
|
2848
|
+
@over
|
|
2849
|
+
def sub(a: Transformation[Scalar], b: Transformation[Scalar]) -> Transformation[Scalar]:
|
|
2850
|
+
""" """
|
|
2851
|
+
...
|
|
2852
|
+
|
|
2853
|
+
@over
|
|
2854
|
+
def sub(a: Tile[Any, Tuple[int, ...]], b: Tile[Any, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
|
|
2855
|
+
"""Subtract each element b from a"""
|
|
2856
|
+
...
|
|
2857
|
+
|
|
2858
|
+
@over
|
|
2859
|
+
def bit_and(a: Int, b: Int) -> Int:
|
|
2860
|
+
""" """
|
|
2861
|
+
...
|
|
2862
|
+
|
|
2863
|
+
@over
|
|
2864
|
+
def bit_or(a: Int, b: Int) -> Int:
|
|
2865
|
+
""" """
|
|
2866
|
+
...
|
|
2867
|
+
|
|
2868
|
+
@over
|
|
2869
|
+
def bit_xor(a: Int, b: Int) -> Int:
|
|
2870
|
+
""" """
|
|
2871
|
+
...
|
|
2872
|
+
|
|
2873
|
+
@over
|
|
2874
|
+
def lshift(a: Int, b: Int) -> Int:
|
|
2875
|
+
""" """
|
|
2876
|
+
...
|
|
2877
|
+
|
|
2878
|
+
@over
|
|
2879
|
+
def rshift(a: Int, b: Int) -> Int:
|
|
2880
|
+
""" """
|
|
2881
|
+
...
|
|
2882
|
+
|
|
2883
|
+
@over
|
|
2884
|
+
def invert(a: Int) -> Int:
|
|
2885
|
+
""" """
|
|
2886
|
+
...
|
|
2887
|
+
|
|
2888
|
+
@over
|
|
2889
|
+
def mul(a: Scalar, b: Scalar) -> Scalar:
|
|
2890
|
+
""" """
|
|
2891
|
+
...
|
|
2892
|
+
|
|
2893
|
+
@over
|
|
2894
|
+
def mul(a: Vector[Any, Scalar], b: Scalar) -> Vector[Any, Scalar]:
|
|
2895
|
+
""" """
|
|
2896
|
+
...
|
|
2897
|
+
|
|
2898
|
+
@over
|
|
2899
|
+
def mul(a: Scalar, b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2900
|
+
""" """
|
|
2901
|
+
...
|
|
2902
|
+
|
|
2903
|
+
@over
|
|
2904
|
+
def mul(a: Quaternion[Scalar], b: Scalar) -> Quaternion[Scalar]:
|
|
2905
|
+
""" """
|
|
2906
|
+
...
|
|
2907
|
+
|
|
2908
|
+
@over
|
|
2909
|
+
def mul(a: Scalar, b: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
2910
|
+
""" """
|
|
2911
|
+
...
|
|
2912
|
+
|
|
2913
|
+
@over
|
|
2914
|
+
def mul(a: Quaternion[Scalar], b: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
2915
|
+
""" """
|
|
2916
|
+
...
|
|
2917
|
+
|
|
2918
|
+
@over
|
|
2919
|
+
def mul(a: Scalar, b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
2920
|
+
""" """
|
|
2921
|
+
...
|
|
2922
|
+
|
|
2923
|
+
@over
|
|
2924
|
+
def mul(a: Matrix[Any, Any, Scalar], b: Scalar) -> Matrix[Any, Any, Scalar]:
|
|
2925
|
+
""" """
|
|
2926
|
+
...
|
|
2927
|
+
|
|
2928
|
+
@over
|
|
2929
|
+
def mul(a: Matrix[Any, Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2930
|
+
""" """
|
|
2931
|
+
...
|
|
2932
|
+
|
|
2933
|
+
@over
|
|
2934
|
+
def mul(a: Vector[Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2935
|
+
""" """
|
|
2936
|
+
...
|
|
2937
|
+
|
|
2938
|
+
@over
|
|
2939
|
+
def mul(a: Matrix[Any, Any, Scalar], b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
2940
|
+
""" """
|
|
2941
|
+
...
|
|
2942
|
+
|
|
2943
|
+
@over
|
|
2944
|
+
def mul(a: Transformation[Scalar], b: Transformation[Scalar]) -> Transformation[Scalar]:
|
|
2945
|
+
""" """
|
|
2946
|
+
...
|
|
2947
|
+
|
|
2948
|
+
@over
|
|
2949
|
+
def mul(a: Scalar, b: Transformation[Scalar]) -> Transformation[Scalar]:
|
|
2950
|
+
""" """
|
|
2951
|
+
...
|
|
2952
|
+
|
|
2953
|
+
@over
|
|
2954
|
+
def mul(a: Transformation[Scalar], b: Scalar) -> Transformation[Scalar]:
|
|
2955
|
+
""" """
|
|
2956
|
+
...
|
|
2957
|
+
|
|
2958
|
+
@over
|
|
2959
|
+
def mul(x: Tile[Any, Tuple[int, ...]], y: Scalar) -> Tile[Any, Tuple[int, ...]]:
|
|
2960
|
+
"""Multiply each element of a tile by a scalar"""
|
|
2961
|
+
...
|
|
2962
|
+
|
|
2963
|
+
@over
|
|
2964
|
+
def mul(x: Scalar, y: Tile[Any, Tuple[int, ...]]) -> Tile[Any, Tuple[int, ...]]:
|
|
2965
|
+
"""Multiply each element of a tile by a scalar"""
|
|
2966
|
+
...
|
|
2967
|
+
|
|
2968
|
+
@over
|
|
2969
|
+
def mod(a: Scalar, b: Scalar) -> Scalar:
|
|
2970
|
+
"""Modulo operation using truncated division."""
|
|
2971
|
+
...
|
|
2972
|
+
|
|
2973
|
+
@over
|
|
2974
|
+
def mod(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Scalar:
|
|
2975
|
+
"""Modulo operation using truncated division."""
|
|
2976
|
+
...
|
|
2977
|
+
|
|
2978
|
+
@over
|
|
2979
|
+
def div(a: Scalar, b: Scalar) -> Scalar:
|
|
2980
|
+
""" """
|
|
2981
|
+
...
|
|
2982
|
+
|
|
2983
|
+
@over
|
|
2984
|
+
def div(a: Vector[Any, Scalar], b: Scalar) -> Vector[Any, Scalar]:
|
|
2985
|
+
""" """
|
|
2986
|
+
...
|
|
2987
|
+
|
|
2988
|
+
@over
|
|
2989
|
+
def div(a: Scalar, b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2990
|
+
""" """
|
|
2991
|
+
...
|
|
2992
|
+
|
|
2993
|
+
@over
|
|
2994
|
+
def div(a: Matrix[Any, Any, Scalar], b: Scalar) -> Matrix[Any, Any, Scalar]:
|
|
2995
|
+
""" """
|
|
2996
|
+
...
|
|
2997
|
+
|
|
2998
|
+
@over
|
|
2999
|
+
def div(a: Scalar, b: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
3000
|
+
""" """
|
|
3001
|
+
...
|
|
3002
|
+
|
|
3003
|
+
@over
|
|
3004
|
+
def div(a: Quaternion[Scalar], b: Scalar) -> Quaternion[Scalar]:
|
|
3005
|
+
""" """
|
|
3006
|
+
...
|
|
3007
|
+
|
|
3008
|
+
@over
|
|
3009
|
+
def div(a: Scalar, b: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
3010
|
+
""" """
|
|
3011
|
+
...
|
|
3012
|
+
|
|
3013
|
+
@over
|
|
3014
|
+
def floordiv(a: Scalar, b: Scalar) -> Scalar:
|
|
3015
|
+
""" """
|
|
3016
|
+
...
|
|
3017
|
+
|
|
3018
|
+
@over
|
|
3019
|
+
def pos(x: Scalar) -> Scalar:
|
|
3020
|
+
""" """
|
|
3021
|
+
...
|
|
3022
|
+
|
|
3023
|
+
@over
|
|
3024
|
+
def pos(x: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
3025
|
+
""" """
|
|
3026
|
+
...
|
|
3027
|
+
|
|
3028
|
+
@over
|
|
3029
|
+
def pos(x: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
3030
|
+
""" """
|
|
3031
|
+
...
|
|
3032
|
+
|
|
3033
|
+
@over
|
|
3034
|
+
def pos(x: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
3035
|
+
""" """
|
|
3036
|
+
...
|
|
3037
|
+
|
|
3038
|
+
@over
|
|
3039
|
+
def neg(x: Scalar) -> Scalar:
|
|
3040
|
+
""" """
|
|
3041
|
+
...
|
|
3042
|
+
|
|
3043
|
+
@over
|
|
3044
|
+
def neg(x: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
3045
|
+
""" """
|
|
3046
|
+
...
|
|
3047
|
+
|
|
3048
|
+
@over
|
|
3049
|
+
def neg(x: Quaternion[Scalar]) -> Quaternion[Scalar]:
|
|
3050
|
+
""" """
|
|
3051
|
+
...
|
|
3052
|
+
|
|
3053
|
+
@over
|
|
3054
|
+
def neg(x: Matrix[Any, Any, Scalar]) -> Matrix[Any, Any, Scalar]:
|
|
3055
|
+
""" """
|
|
3056
|
+
...
|
|
3057
|
+
|
|
3058
|
+
@over
|
|
3059
|
+
def neg(x: Tile[Any, Tuple[int, ...]]) -> Tile[Scalar, Tuple[int, ...]]:
|
|
3060
|
+
"""Negate each element of a tile"""
|
|
3061
|
+
...
|
|
3062
|
+
|
|
3063
|
+
@over
|
|
3064
|
+
def unot(a: bool) -> bool:
|
|
3065
|
+
""" """
|
|
3066
|
+
...
|
|
3067
|
+
|
|
3068
|
+
@over
|
|
3069
|
+
def unot(a: int8) -> bool:
|
|
3070
|
+
""" """
|
|
3071
|
+
...
|
|
3072
|
+
|
|
3073
|
+
@over
|
|
3074
|
+
def unot(a: uint8) -> bool:
|
|
3075
|
+
""" """
|
|
3076
|
+
...
|
|
3077
|
+
|
|
3078
|
+
@over
|
|
3079
|
+
def unot(a: int16) -> bool:
|
|
3080
|
+
""" """
|
|
3081
|
+
...
|
|
3082
|
+
|
|
3083
|
+
@over
|
|
3084
|
+
def unot(a: uint16) -> bool:
|
|
3085
|
+
""" """
|
|
3086
|
+
...
|
|
3087
|
+
|
|
3088
|
+
@over
|
|
3089
|
+
def unot(a: int32) -> bool:
|
|
3090
|
+
""" """
|
|
3091
|
+
...
|
|
3092
|
+
|
|
3093
|
+
@over
|
|
3094
|
+
def unot(a: uint32) -> bool:
|
|
3095
|
+
""" """
|
|
3096
|
+
...
|
|
3097
|
+
|
|
3098
|
+
@over
|
|
3099
|
+
def unot(a: int64) -> bool:
|
|
3100
|
+
""" """
|
|
3101
|
+
...
|
|
3102
|
+
|
|
3103
|
+
@over
|
|
3104
|
+
def unot(a: uint64) -> bool:
|
|
3105
|
+
""" """
|
|
3106
|
+
...
|
|
3107
|
+
|
|
3108
|
+
@over
|
|
3109
|
+
def unot(a: Array[Any]) -> bool:
|
|
3110
|
+
""" """
|
|
3111
|
+
...
|
|
3112
|
+
|
|
3113
|
+
@over
|
|
3114
|
+
def tile_diag_add(a: Tile[Any, Tuple[int, int]], d: Tile[Any, Tuple[int]]) -> Tile[Any, Tuple[int, int]]:
|
|
3115
|
+
"""Add a square matrix and a diagonal matrix 'd' represented as a 1D tile"""
|
|
3116
|
+
...
|
|
3117
|
+
|
|
3118
|
+
@over
|
|
3119
|
+
def tile_matmul(a: Tile[Float, Tuple[int, int]], b: Tile[Float, Tuple[int, int]], out: Tile[Float, Tuple[int, int]]):
|
|
3120
|
+
"""Computes the matrix product and accumulates ``out += a*b``.
|
|
3121
|
+
|
|
3122
|
+
Supported datatypes are:
|
|
3123
|
+
* fp16, fp32, fp64 (real)
|
|
3124
|
+
* vec2h, vec2f, vec2d (complex)
|
|
3125
|
+
|
|
3126
|
+
All input and output tiles must have the same datatype. Tile data will automatically be migrated
|
|
3127
|
+
to shared memory if necessary and will use TensorCore operations when available.
|
|
3128
|
+
|
|
3129
|
+
:param a: A tile with ``shape=(M, K)``
|
|
3130
|
+
:param b: A tile with ``shape=(K, N)``
|
|
3131
|
+
:param out: A tile with ``shape=(M, N)``
|
|
3132
|
+
|
|
3133
|
+
"""
|
|
3134
|
+
...
|
|
3135
|
+
|
|
3136
|
+
@over
|
|
3137
|
+
def tile_matmul(a: Tile[Float, Tuple[int, int]], b: Tile[Float, Tuple[int, int]]) -> Tile[Float, Tuple[int, int]]:
|
|
3138
|
+
"""Computes the matrix product ``out = a*b``.
|
|
3139
|
+
|
|
3140
|
+
Supported datatypes are:
|
|
3141
|
+
* fp16, fp32, fp64 (real)
|
|
3142
|
+
* vec2h, vec2f, vec2d (complex)
|
|
3143
|
+
|
|
3144
|
+
Both input tiles must have the same datatype. Tile data will automatically be migrated
|
|
3145
|
+
to shared memory if necessary and will use TensorCore operations when available.
|
|
3146
|
+
|
|
3147
|
+
:param a: A tile with ``shape=(M, K)``
|
|
3148
|
+
:param b: A tile with ``shape=(K, N)``
|
|
3149
|
+
:returns: A tile with ``shape=(M, N)``
|
|
3150
|
+
|
|
3151
|
+
"""
|
|
3152
|
+
...
|
|
3153
|
+
|
|
3154
|
+
@over
|
|
3155
|
+
def tile_fft(inout: Tile[Vector[2, Float], Tuple[int, int]]) -> Tile[Vector[2, Float], Tuple[int, int]]:
|
|
3156
|
+
"""Compute the forward FFT along the second dimension of a 2D tile of data.
|
|
3157
|
+
|
|
3158
|
+
This function cooperatively computes the forward FFT on a tile of data inplace, treating each row individually.
|
|
3159
|
+
|
|
3160
|
+
Note that computing the adjoint is not yet supported.
|
|
3161
|
+
|
|
3162
|
+
Supported datatypes are:
|
|
3163
|
+
* vec2f, vec2d
|
|
3164
|
+
|
|
3165
|
+
:param inout: The input/output tile
|
|
3166
|
+
"""
|
|
3167
|
+
...
|
|
3168
|
+
|
|
3169
|
+
@over
|
|
3170
|
+
def tile_ifft(inout: Tile[Vector[2, Float], Tuple[int, int]]) -> Tile[Vector[2, Float], Tuple[int, int]]:
|
|
3171
|
+
"""Compute the inverse FFT along the second dimension of a 2D tile of data.
|
|
3172
|
+
|
|
3173
|
+
This function cooperatively computes the inverse FFT on a tile of data inplace, treating each row individually.
|
|
3174
|
+
|
|
3175
|
+
Note that computing the adjoint is not yet supported.
|
|
3176
|
+
|
|
3177
|
+
Supported datatypes are:
|
|
3178
|
+
* vec2f, vec2d
|
|
3179
|
+
|
|
3180
|
+
:param inout: The input/output tile
|
|
3181
|
+
"""
|
|
3182
|
+
...
|
|
3183
|
+
|
|
3184
|
+
@over
|
|
3185
|
+
def tile_cholesky(A: Tile[Float, Tuple[int, int]]) -> Tile[Float, Tuple[int, int]]:
|
|
3186
|
+
"""Compute the Cholesky factorization L of a matrix A.
|
|
3187
|
+
L is lower triangular and satisfies LL^T = A.
|
|
3188
|
+
|
|
3189
|
+
Only the lower triangular portion of A is used for the decomposition;
|
|
3190
|
+
the upper triangular part may be left unspecified.
|
|
3191
|
+
|
|
3192
|
+
Note that computing the adjoint is not yet supported.
|
|
3193
|
+
|
|
3194
|
+
Supported datatypes are:
|
|
3195
|
+
* float32
|
|
3196
|
+
* float64
|
|
3197
|
+
|
|
3198
|
+
:param A: A square, symmetric positive-definite, matrix. Only the lower triangular part of A is needed; the upper part is ignored.
|
|
3199
|
+
:returns L: A square, lower triangular, matrix, such that LL^T = A
|
|
3200
|
+
"""
|
|
3201
|
+
...
|
|
3202
|
+
|
|
3203
|
+
@over
|
|
3204
|
+
def tile_cholesky_solve(L: Tile[Float, Tuple[int, int]], y: Tile[Float, Tuple[int]]):
|
|
3205
|
+
"""With L such that LL^T = A, solve for x in Ax = y
|
|
3206
|
+
|
|
3207
|
+
Note that computing the adjoint is not yet supported.
|
|
3208
|
+
|
|
3209
|
+
Supported datatypes are:
|
|
3210
|
+
* float32
|
|
3211
|
+
* float64
|
|
3212
|
+
|
|
3213
|
+
:param L: A square, lower triangular, matrix, such that LL^T = A
|
|
3214
|
+
:param y: A 1D or 2D tile of length M
|
|
3215
|
+
:returns x: A tile of the same shape as y such that LL^T x = y
|
|
3216
|
+
"""
|
|
3217
|
+
...
|
|
3218
|
+
|
|
3219
|
+
@over
|
|
3220
|
+
def tile_lower_solve(L: Tile[Float, Tuple[int, int]], y: Tile[Float, Tuple[int]]) -> Tile[Float, Tuple[int]]:
|
|
3221
|
+
"""Solve for z in Lz = y, where L is a lower triangular matrix.
|
|
3222
|
+
|
|
3223
|
+
This performs general forward substitution for a lower triangular system.
|
|
3224
|
+
|
|
3225
|
+
Note that computing the adjoint is not yet supported.
|
|
3226
|
+
|
|
3227
|
+
Supported datatypes are:
|
|
3228
|
+
* float32
|
|
3229
|
+
* float64
|
|
3230
|
+
|
|
3231
|
+
:param L: A square, non-singular, lower triangular matrix
|
|
3232
|
+
:param y: A 1D or 2D tile with compatible shape
|
|
3233
|
+
:returns z: A tile of the same shape as y such that Lz = y
|
|
3234
|
+
"""
|
|
3235
|
+
...
|
|
3236
|
+
|
|
3237
|
+
@over
|
|
3238
|
+
def tile_upper_solve(U: Tile[Float, Tuple[int, int]], z: Tile[Float, Tuple[int]]) -> Tile[Float, Tuple[int]]:
|
|
3239
|
+
"""Solve for x in U x = z, where U is an upper triangular matrix.
|
|
3240
|
+
|
|
3241
|
+
This performs general back substitution for upper triangular systems.
|
|
3242
|
+
|
|
3243
|
+
Note that computing the adjoint is not yet supported.
|
|
3244
|
+
|
|
3245
|
+
Supported datatypes are:
|
|
3246
|
+
* float32
|
|
3247
|
+
* float64
|
|
3248
|
+
|
|
3249
|
+
:param U: A square, non-singular, upper triangular matrix
|
|
3250
|
+
:param z: A 1D or 2D tile with compatible shape
|
|
3251
|
+
:returns x: A tile of the same shape as z such that U x = z
|
|
3252
|
+
"""
|
|
3253
|
+
...
|
|
3254
|
+
|
|
3255
|
+
@over
|
|
3256
|
+
def static(expr: Any) -> Any:
|
|
3257
|
+
"""Evaluate a static Python expression and replaces it with its result.
|
|
3258
|
+
|
|
3259
|
+
See the :ref:`code generation guide <static_expressions>` for more details.
|
|
3260
|
+
|
|
3261
|
+
The inner expression must only reference variables that are available from the current scope where the Warp kernel or function containing the expression is defined,
|
|
3262
|
+
which includes constant variables and variables captured in the current closure in which the function or kernel is implemented.
|
|
3263
|
+
The return type of the expression must be either a Warp function, a string, or a type that is supported inside Warp kernels and functions
|
|
3264
|
+
(excluding Warp arrays since they cannot be created in a Warp kernel at the moment).
|
|
3265
|
+
"""
|
|
3266
|
+
...
|
|
3267
|
+
|
|
3268
|
+
@over
|
|
3269
|
+
def len(a: Vector[Any, Scalar]) -> int:
|
|
3270
|
+
"""Return the number of elements in a vector."""
|
|
3271
|
+
...
|
|
3272
|
+
|
|
3273
|
+
@over
|
|
3274
|
+
def len(a: Quaternion[Scalar]) -> int:
|
|
3275
|
+
"""Return the number of elements in a quaternion."""
|
|
3276
|
+
...
|
|
3277
|
+
|
|
3278
|
+
@over
|
|
3279
|
+
def len(a: Matrix[Any, Any, Scalar]) -> int:
|
|
3280
|
+
"""Return the number of rows in a matrix."""
|
|
3281
|
+
...
|
|
3282
|
+
|
|
3283
|
+
@over
|
|
3284
|
+
def len(a: Transformation[Float]) -> int:
|
|
3285
|
+
"""Return the number of elements in a transformation."""
|
|
3286
|
+
...
|
|
3287
|
+
|
|
3288
|
+
@over
|
|
3289
|
+
def len(a: Array[Any]) -> int:
|
|
3290
|
+
"""Return the size of the first dimension in an array."""
|
|
3291
|
+
...
|
|
3292
|
+
|
|
3293
|
+
@over
|
|
3294
|
+
def len(a: Tile[Any, Tuple[int, ...]]) -> int:
|
|
3295
|
+
"""Return the number of rows in a tile."""
|
|
3296
|
+
...
|
|
3297
|
+
|
|
3298
|
+
@over
|
|
3299
|
+
def len(a: Tuple) -> int:
|
|
3300
|
+
"""Return the number of elements in a tuple."""
|
|
3301
|
+
...
|
|
3302
|
+
|
|
3303
|
+
@over
|
|
3304
|
+
def norm_l1(v: Any):
|
|
3305
|
+
"""Computes the L1 norm of a vector v.
|
|
3306
|
+
|
|
3307
|
+
.. math:: \|v\|_1 = \sum_i |v_i|
|
|
3308
|
+
|
|
3309
|
+
Args:
|
|
3310
|
+
v (Vector[Any,Float]): The vector to compute the L1 norm of.
|
|
3311
|
+
|
|
3312
|
+
Returns:
|
|
3313
|
+
float: The L1 norm of the vector.
|
|
3314
|
+
"""
|
|
3315
|
+
...
|
|
3316
|
+
|
|
3317
|
+
@over
|
|
3318
|
+
def norm_l2(v: Any):
|
|
3319
|
+
"""Computes the L2 norm of a vector v.
|
|
3320
|
+
|
|
3321
|
+
.. math:: \|v\|_2 = \sqrt{\sum_i v_i^2}
|
|
3322
|
+
|
|
3323
|
+
Args:
|
|
3324
|
+
v (Vector[Any,Float]): The vector to compute the L2 norm of.
|
|
3325
|
+
|
|
3326
|
+
Returns:
|
|
3327
|
+
float: The L2 norm of the vector.
|
|
3328
|
+
"""
|
|
3329
|
+
...
|
|
3330
|
+
|
|
3331
|
+
@over
|
|
3332
|
+
def norm_huber(v: Any, delta: float):
|
|
3333
|
+
"""Computes the Huber norm of a vector v with a given delta.
|
|
3334
|
+
|
|
3335
|
+
.. math::
|
|
3336
|
+
H(v) = \begin{cases} \frac{1}{2} \|v\|^2 & \text{if } \|v\| \leq \delta \\ \delta(\|v\| - \frac{1}{2}\delta) & \text{otherwise} \end{cases}
|
|
3337
|
+
|
|
3338
|
+
.. image:: /img/norm_huber.svg
|
|
3339
|
+
:align: center
|
|
3340
|
+
|
|
3341
|
+
Args:
|
|
3342
|
+
v (Vector[Any,Float]): The vector to compute the Huber norm of.
|
|
3343
|
+
delta (float): The threshold value, defaults to 1.0.
|
|
3344
|
+
|
|
3345
|
+
Returns:
|
|
3346
|
+
float: The Huber norm of the vector.
|
|
3347
|
+
"""
|
|
3348
|
+
...
|
|
3349
|
+
|
|
3350
|
+
@over
|
|
3351
|
+
def norm_pseudo_huber(v: Any, delta: float):
|
|
3352
|
+
"""Computes the "pseudo" Huber norm of a vector v with a given delta.
|
|
3353
|
+
|
|
3354
|
+
.. math::
|
|
3355
|
+
H^\prime(v) = \delta \sqrt{1 + \frac{\|v\|^2}{\delta^2}}
|
|
3356
|
+
|
|
3357
|
+
.. image:: /img/norm_pseudo_huber.svg
|
|
3358
|
+
:align: center
|
|
3359
|
+
|
|
3360
|
+
Args:
|
|
3361
|
+
v (Vector[Any,Float]): The vector to compute the Huber norm of.
|
|
3362
|
+
delta (float): The threshold value, defaults to 1.0.
|
|
3363
|
+
|
|
3364
|
+
Returns:
|
|
3365
|
+
float: The Huber norm of the vector.
|
|
3366
|
+
"""
|
|
3367
|
+
...
|
|
3368
|
+
|
|
3369
|
+
@over
|
|
3370
|
+
def smooth_normalize(v: Any, delta: float):
|
|
3371
|
+
"""Normalizes a vector using the pseudo-Huber norm.
|
|
3372
|
+
|
|
3373
|
+
See :func:`norm_pseudo_huber`.
|
|
3374
|
+
|
|
3375
|
+
.. math::
|
|
3376
|
+
\frac{v}{H^\prime(v)}
|
|
3377
|
+
|
|
3378
|
+
Args:
|
|
3379
|
+
v (Vector[Any,Float]): The vector to normalize.
|
|
3380
|
+
delta (float): The threshold value, defaults to 1.0.
|
|
3381
|
+
|
|
3382
|
+
Returns:
|
|
3383
|
+
Vector[Any,Float]: The normalized vector.
|
|
3384
|
+
"""
|
|
3385
|
+
...
|
|
3386
|
+
|
|
3387
|
+
@over
|
|
3388
|
+
def transform_from_matrix(mat: Matrix[4, 4, float32]) -> Transformation[float32]:
|
|
3389
|
+
"""Construct a transformation from a 4x4 matrix.
|
|
3390
|
+
|
|
3391
|
+
.. math::
|
|
3392
|
+
M = \begin{bmatrix}
|
|
3393
|
+
R_{00} & R_{01} & R_{02} & p_x \\
|
|
3394
|
+
R_{10} & R_{11} & R_{12} & p_y \\
|
|
3395
|
+
R_{20} & R_{21} & R_{22} & p_z \\
|
|
3396
|
+
0 & 0 & 0 & 1
|
|
3397
|
+
\end{bmatrix}
|
|
3398
|
+
|
|
3399
|
+
Where:
|
|
3400
|
+
|
|
3401
|
+
* :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
|
|
3402
|
+
* :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
|
|
3403
|
+
|
|
3404
|
+
Args:
|
|
3405
|
+
mat (Matrix[4, 4, Float]): Matrix to convert.
|
|
3406
|
+
|
|
3407
|
+
Returns:
|
|
3408
|
+
Transformation[Float]: The transformation.
|
|
3409
|
+
"""
|
|
3410
|
+
...
|
|
3411
|
+
|
|
3412
|
+
@over
|
|
3413
|
+
def transform_to_matrix(xform: Transformation[float32]) -> Matrix[4, 4, float32]:
|
|
3414
|
+
"""Convert a transformation to a 4x4 matrix.
|
|
3415
|
+
|
|
3416
|
+
.. math::
|
|
3417
|
+
M = \begin{bmatrix}
|
|
3418
|
+
R_{00} & R_{01} & R_{02} & p_x \\
|
|
3419
|
+
R_{10} & R_{11} & R_{12} & p_y \\
|
|
3420
|
+
R_{20} & R_{21} & R_{22} & p_z \\
|
|
3421
|
+
0 & 0 & 0 & 1
|
|
3422
|
+
\end{bmatrix}
|
|
3423
|
+
|
|
3424
|
+
Where:
|
|
3425
|
+
|
|
3426
|
+
* :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
|
|
3427
|
+
* :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
|
|
3428
|
+
|
|
3429
|
+
Args:
|
|
3430
|
+
xform (Transformation[Float]): Transformation to convert.
|
|
3431
|
+
|
|
3432
|
+
Returns:
|
|
3433
|
+
Matrix[4, 4, Float]: The matrix.
|
|
3434
|
+
"""
|
|
3435
|
+
...
|
|
3436
|
+
|
|
3437
|
+
@over
|
|
3438
|
+
def transform_compose(position: Vector[3, float32], rotation: Quaternion[float32], scale: Vector[3, float32]):
|
|
3439
|
+
"""Compose a 4x4 transformation matrix from a 3D position, quaternion orientation, and 3D scale.
|
|
3440
|
+
|
|
3441
|
+
.. math::
|
|
3442
|
+
M = \begin{bmatrix}
|
|
3443
|
+
s_x R_{00} & s_y R_{01} & s_z R_{02} & p_x \\
|
|
3444
|
+
s_x R_{10} & s_y R_{11} & s_z R_{12} & p_y \\
|
|
3445
|
+
s_x R_{20} & s_y R_{21} & s_z R_{22} & p_z \\
|
|
3446
|
+
0 & 0 & 0 & 1
|
|
3447
|
+
\end{bmatrix}
|
|
3448
|
+
|
|
3449
|
+
Where:
|
|
3450
|
+
|
|
3451
|
+
* :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
|
|
3452
|
+
* :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
|
|
3453
|
+
* :math:`s` is the 3D scale vector :math:`[s_x, s_y, s_z]` of the input transform.
|
|
3454
|
+
|
|
3455
|
+
Args:
|
|
3456
|
+
position (Vector[3, Float]): The 3D position vector.
|
|
3457
|
+
rotation (Quaternion[Float]): The quaternion orientation.
|
|
3458
|
+
scale (Vector[3, Float]): The 3D scale vector.
|
|
3459
|
+
|
|
3460
|
+
Returns:
|
|
3461
|
+
Matrix[4, 4, Float]: The transformation matrix.
|
|
3462
|
+
"""
|
|
3463
|
+
...
|
|
3464
|
+
|
|
3465
|
+
@over
|
|
3466
|
+
def transform_decompose(m: Matrix[4, 4, float32]):
|
|
3467
|
+
"""Decompose a 4x4 transformation matrix into 3D position, quaternion orientation, and 3D scale.
|
|
3468
|
+
|
|
3469
|
+
.. math::
|
|
3470
|
+
M = \begin{bmatrix}
|
|
3471
|
+
s_x R_{00} & s_y R_{01} & s_z R_{02} & p_x \\
|
|
3472
|
+
s_x R_{10} & s_y R_{11} & s_z R_{12} & p_y \\
|
|
3473
|
+
s_x R_{20} & s_y R_{21} & s_z R_{22} & p_z \\
|
|
3474
|
+
0 & 0 & 0 & 1
|
|
3475
|
+
\end{bmatrix}
|
|
3476
|
+
|
|
3477
|
+
Where:
|
|
3478
|
+
|
|
3479
|
+
* :math:`R` is the 3x3 rotation matrix created from the orientation quaternion of the input transform.
|
|
3480
|
+
* :math:`p` is the 3D position vector :math:`[p_x, p_y, p_z]` of the input transform.
|
|
3481
|
+
* :math:`s` is the 3D scale vector :math:`[s_x, s_y, s_z]` of the input transform.
|
|
3482
|
+
|
|
3483
|
+
Args:
|
|
3484
|
+
m (Matrix[4, 4, Float]): The matrix to decompose.
|
|
3485
|
+
|
|
3486
|
+
Returns:
|
|
3487
|
+
Tuple[Vector[3, Float], Quaternion[Float], Vector[3, Float]]: A tuple containing the position vector, quaternion orientation, and scale vector.
|
|
3488
|
+
"""
|
|
3489
|
+
...
|