warp-lang 1.8.1__py3-none-win_amd64.whl → 1.9.1__py3-none-win_amd64.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 +282 -103
- warp/__init__.pyi +1904 -114
- warp/bin/warp-clang.dll +0 -0
- warp/bin/warp.dll +0 -0
- warp/build.py +93 -30
- warp/build_dll.py +331 -101
- warp/builtins.py +1244 -160
- warp/codegen.py +317 -206
- warp/config.py +1 -1
- warp/context.py +1465 -789
- warp/examples/core/example_marching_cubes.py +1 -0
- warp/examples/core/example_render_opengl.py +100 -3
- warp/examples/fem/example_apic_fluid.py +98 -52
- warp/examples/fem/example_convection_diffusion_dg.py +25 -4
- warp/examples/fem/example_diffusion_mgpu.py +8 -3
- warp/examples/fem/utils.py +68 -22
- warp/examples/interop/example_jax_kernel.py +2 -1
- warp/fabric.py +1 -1
- warp/fem/cache.py +27 -19
- warp/fem/domain.py +2 -2
- warp/fem/field/nodal_field.py +2 -2
- warp/fem/field/virtual.py +264 -166
- warp/fem/geometry/geometry.py +5 -5
- warp/fem/integrate.py +129 -51
- warp/fem/space/restriction.py +4 -0
- warp/fem/space/shape/tet_shape_function.py +3 -10
- warp/jax_experimental/custom_call.py +25 -2
- warp/jax_experimental/ffi.py +22 -1
- warp/jax_experimental/xla_ffi.py +16 -7
- warp/marching_cubes.py +708 -0
- warp/native/array.h +99 -4
- warp/native/builtin.h +86 -9
- warp/native/bvh.cpp +64 -28
- warp/native/bvh.cu +58 -58
- warp/native/bvh.h +2 -2
- warp/native/clang/clang.cpp +7 -7
- warp/native/coloring.cpp +8 -2
- warp/native/crt.cpp +2 -2
- warp/native/crt.h +3 -5
- warp/native/cuda_util.cpp +41 -10
- warp/native/cuda_util.h +10 -4
- warp/native/exports.h +1842 -1908
- warp/native/fabric.h +2 -1
- warp/native/hashgrid.cpp +37 -37
- warp/native/hashgrid.cu +2 -2
- warp/native/initializer_array.h +1 -1
- warp/native/intersect.h +2 -2
- warp/native/mat.h +1910 -116
- warp/native/mathdx.cpp +43 -43
- warp/native/mesh.cpp +24 -24
- warp/native/mesh.cu +26 -26
- warp/native/mesh.h +4 -2
- warp/native/nanovdb/GridHandle.h +179 -12
- warp/native/nanovdb/HostBuffer.h +8 -7
- warp/native/nanovdb/NanoVDB.h +517 -895
- warp/native/nanovdb/NodeManager.h +323 -0
- warp/native/nanovdb/PNanoVDB.h +2 -2
- warp/native/quat.h +331 -14
- warp/native/range.h +7 -1
- warp/native/reduce.cpp +10 -10
- warp/native/reduce.cu +13 -14
- warp/native/runlength_encode.cpp +2 -2
- warp/native/runlength_encode.cu +5 -5
- warp/native/scan.cpp +3 -3
- warp/native/scan.cu +4 -4
- warp/native/sort.cpp +10 -10
- warp/native/sort.cu +40 -31
- warp/native/sort.h +2 -0
- warp/native/sparse.cpp +8 -8
- warp/native/sparse.cu +13 -13
- warp/native/spatial.h +366 -17
- warp/native/temp_buffer.h +2 -2
- warp/native/tile.h +471 -82
- warp/native/vec.h +328 -14
- warp/native/volume.cpp +54 -54
- warp/native/volume.cu +1 -1
- warp/native/volume.h +2 -1
- warp/native/volume_builder.cu +30 -37
- warp/native/warp.cpp +150 -149
- warp/native/warp.cu +377 -216
- warp/native/warp.h +227 -226
- warp/optim/linear.py +736 -271
- warp/render/imgui_manager.py +289 -0
- warp/render/render_opengl.py +99 -18
- warp/render/render_usd.py +1 -0
- warp/sim/graph_coloring.py +2 -2
- warp/sparse.py +558 -175
- warp/tests/aux_test_module_aot.py +7 -0
- warp/tests/cuda/test_async.py +3 -3
- warp/tests/cuda/test_conditional_captures.py +101 -0
- warp/tests/geometry/test_hash_grid.py +38 -0
- warp/tests/geometry/test_marching_cubes.py +233 -12
- warp/tests/interop/test_jax.py +608 -28
- warp/tests/sim/test_coloring.py +6 -6
- warp/tests/test_array.py +58 -5
- warp/tests/test_codegen.py +4 -3
- warp/tests/test_context.py +8 -15
- warp/tests/test_enum.py +136 -0
- warp/tests/test_examples.py +2 -2
- warp/tests/test_fem.py +49 -6
- warp/tests/test_fixedarray.py +229 -0
- warp/tests/test_func.py +18 -15
- warp/tests/test_future_annotations.py +7 -5
- warp/tests/test_linear_solvers.py +30 -0
- warp/tests/test_map.py +15 -1
- warp/tests/test_mat.py +1518 -378
- warp/tests/test_mat_assign_copy.py +178 -0
- warp/tests/test_mat_constructors.py +574 -0
- warp/tests/test_module_aot.py +287 -0
- warp/tests/test_print.py +69 -0
- warp/tests/test_quat.py +140 -34
- warp/tests/test_quat_assign_copy.py +145 -0
- warp/tests/test_reload.py +2 -1
- warp/tests/test_sparse.py +71 -0
- warp/tests/test_spatial.py +140 -34
- warp/tests/test_spatial_assign_copy.py +160 -0
- warp/tests/test_struct.py +43 -3
- warp/tests/test_tuple.py +96 -0
- warp/tests/test_types.py +61 -20
- warp/tests/test_vec.py +179 -34
- warp/tests/test_vec_assign_copy.py +143 -0
- warp/tests/tile/test_tile.py +245 -18
- warp/tests/tile/test_tile_cholesky.py +605 -0
- warp/tests/tile/test_tile_load.py +169 -0
- warp/tests/tile/test_tile_mathdx.py +2 -558
- warp/tests/tile/test_tile_matmul.py +1 -1
- warp/tests/tile/test_tile_mlp.py +1 -1
- warp/tests/tile/test_tile_shared_memory.py +5 -5
- warp/tests/unittest_suites.py +6 -0
- warp/tests/walkthrough_debug.py +1 -1
- warp/thirdparty/unittest_parallel.py +108 -9
- warp/types.py +571 -267
- warp/utils.py +68 -86
- {warp_lang-1.8.1.dist-info → warp_lang-1.9.1.dist-info}/METADATA +29 -69
- {warp_lang-1.8.1.dist-info → warp_lang-1.9.1.dist-info}/RECORD +138 -128
- warp/native/marching.cpp +0 -19
- warp/native/marching.cu +0 -514
- warp/native/marching.h +0 -19
- {warp_lang-1.8.1.dist-info → warp_lang-1.9.1.dist-info}/WHEEL +0 -0
- {warp_lang-1.8.1.dist-info → warp_lang-1.9.1.dist-info}/licenses/LICENSE.md +0 -0
- {warp_lang-1.8.1.dist-info → warp_lang-1.9.1.dist-info}/top_level.txt +0 -0
warp/__init__.pyi
CHANGED
|
@@ -20,6 +20,7 @@ from typing import Tuple
|
|
|
20
20
|
from typing import Callable
|
|
21
21
|
from typing import TypeVar
|
|
22
22
|
from typing import Generic
|
|
23
|
+
from typing import Sequence
|
|
23
24
|
from typing import overload as over
|
|
24
25
|
|
|
25
26
|
Length = TypeVar("Length", bound=int)
|
|
@@ -36,122 +37,1718 @@ FabricArray = Generic[DType]
|
|
|
36
37
|
IndexedFabricArray = Generic[DType]
|
|
37
38
|
Tile = Generic[DType, Shape]
|
|
38
39
|
|
|
39
|
-
from warp.types import array
|
|
40
|
-
from warp.types import
|
|
41
|
-
from warp.
|
|
42
|
-
from warp.types import
|
|
40
|
+
from warp.types import array as array
|
|
41
|
+
from warp.types import array1d as array1d
|
|
42
|
+
from warp.types import array2d as array2d
|
|
43
|
+
from warp.types import array3d as array3d
|
|
44
|
+
from warp.types import array4d as array4d
|
|
45
|
+
from warp.types import constant as constant
|
|
46
|
+
from warp.types import from_ptr as from_ptr
|
|
47
|
+
from warp.types import fixedarray as fixedarray
|
|
48
|
+
from warp.types import indexedarray as indexedarray
|
|
49
|
+
from warp.types import indexedarray1d as indexedarray1d
|
|
50
|
+
from warp.types import indexedarray2d as indexedarray2d
|
|
51
|
+
from warp.types import indexedarray3d as indexedarray3d
|
|
52
|
+
from warp.types import indexedarray4d as indexedarray4d
|
|
53
|
+
from warp.fabric import fabricarray as fabricarray
|
|
54
|
+
from warp.fabric import fabricarrayarray as fabricarrayarray
|
|
55
|
+
from warp.fabric import indexedfabricarray as indexedfabricarray
|
|
56
|
+
from warp.fabric import indexedfabricarrayarray as indexedfabricarrayarray
|
|
57
|
+
from warp.types import tile as tile
|
|
58
|
+
|
|
59
|
+
from warp.types import bool as bool
|
|
60
|
+
from warp.types import int8 as int8
|
|
61
|
+
from warp.types import uint8 as uint8
|
|
62
|
+
from warp.types import int16 as int16
|
|
63
|
+
from warp.types import uint16 as uint16
|
|
64
|
+
from warp.types import int32 as int32
|
|
65
|
+
from warp.types import uint32 as uint32
|
|
66
|
+
from warp.types import int64 as int64
|
|
67
|
+
from warp.types import uint64 as uint64
|
|
68
|
+
from warp.types import float16 as float16
|
|
69
|
+
from warp.types import float32 as float32
|
|
70
|
+
from warp.types import float64 as float64
|
|
71
|
+
|
|
72
|
+
from warp.types import vec2 as vec2
|
|
73
|
+
from warp.types import vec2b as vec2b
|
|
74
|
+
from warp.types import vec2ub as vec2ub
|
|
75
|
+
from warp.types import vec2s as vec2s
|
|
76
|
+
from warp.types import vec2us as vec2us
|
|
77
|
+
from warp.types import vec2i as vec2i
|
|
78
|
+
from warp.types import vec2ui as vec2ui
|
|
79
|
+
from warp.types import vec2l as vec2l
|
|
80
|
+
from warp.types import vec2ul as vec2ul
|
|
81
|
+
from warp.types import vec2h as vec2h
|
|
82
|
+
from warp.types import vec2f as vec2f
|
|
83
|
+
from warp.types import vec2d as vec2d
|
|
84
|
+
|
|
85
|
+
from warp.types import vec3 as vec3
|
|
86
|
+
from warp.types import vec3b as vec3b
|
|
87
|
+
from warp.types import vec3ub as vec3ub
|
|
88
|
+
from warp.types import vec3s as vec3s
|
|
89
|
+
from warp.types import vec3us as vec3us
|
|
90
|
+
from warp.types import vec3i as vec3i
|
|
91
|
+
from warp.types import vec3ui as vec3ui
|
|
92
|
+
from warp.types import vec3l as vec3l
|
|
93
|
+
from warp.types import vec3ul as vec3ul
|
|
94
|
+
from warp.types import vec3h as vec3h
|
|
95
|
+
from warp.types import vec3f as vec3f
|
|
96
|
+
from warp.types import vec3d as vec3d
|
|
97
|
+
|
|
98
|
+
from warp.types import vec4 as vec4
|
|
99
|
+
from warp.types import vec4b as vec4b
|
|
100
|
+
from warp.types import vec4ub as vec4ub
|
|
101
|
+
from warp.types import vec4s as vec4s
|
|
102
|
+
from warp.types import vec4us as vec4us
|
|
103
|
+
from warp.types import vec4i as vec4i
|
|
104
|
+
from warp.types import vec4ui as vec4ui
|
|
105
|
+
from warp.types import vec4l as vec4l
|
|
106
|
+
from warp.types import vec4ul as vec4ul
|
|
107
|
+
from warp.types import vec4h as vec4h
|
|
108
|
+
from warp.types import vec4f as vec4f
|
|
109
|
+
from warp.types import vec4d as vec4d
|
|
110
|
+
|
|
111
|
+
from warp.types import mat22 as mat22
|
|
112
|
+
from warp.types import mat22h as mat22h
|
|
113
|
+
from warp.types import mat22f as mat22f
|
|
114
|
+
from warp.types import mat22d as mat22d
|
|
115
|
+
|
|
116
|
+
from warp.types import mat33 as mat33
|
|
117
|
+
from warp.types import mat33h as mat33h
|
|
118
|
+
from warp.types import mat33f as mat33f
|
|
119
|
+
from warp.types import mat33d as mat33d
|
|
120
|
+
|
|
121
|
+
from warp.types import mat44 as mat44
|
|
122
|
+
from warp.types import mat44h as mat44h
|
|
123
|
+
from warp.types import mat44f as mat44f
|
|
124
|
+
from warp.types import mat44d as mat44d
|
|
125
|
+
|
|
126
|
+
from warp.types import quat as quat
|
|
127
|
+
from warp.types import quath as quath
|
|
128
|
+
from warp.types import quatf as quatf
|
|
129
|
+
from warp.types import quatd as quatd
|
|
130
|
+
|
|
131
|
+
from warp.types import transform as transform
|
|
132
|
+
from warp.types import transformh as transformh
|
|
133
|
+
from warp.types import transformf as transformf
|
|
134
|
+
from warp.types import transformd as transformd
|
|
135
|
+
|
|
136
|
+
from warp.types import spatial_vector as spatial_vector
|
|
137
|
+
from warp.types import spatial_vectorh as spatial_vectorh
|
|
138
|
+
from warp.types import spatial_vectorf as spatial_vectorf
|
|
139
|
+
from warp.types import spatial_vectord as spatial_vectord
|
|
140
|
+
|
|
141
|
+
from warp.types import spatial_matrix as spatial_matrix
|
|
142
|
+
from warp.types import spatial_matrixh as spatial_matrixh
|
|
143
|
+
from warp.types import spatial_matrixf as spatial_matrixf
|
|
144
|
+
from warp.types import spatial_matrixd as spatial_matrixd
|
|
145
|
+
|
|
146
|
+
from warp.types import Int as Int
|
|
147
|
+
from warp.types import Float as Float
|
|
148
|
+
from warp.types import Scalar as Scalar
|
|
149
|
+
|
|
150
|
+
from warp.types import Bvh as Bvh
|
|
151
|
+
from warp.types import Mesh as Mesh
|
|
152
|
+
from warp.types import HashGrid as HashGrid
|
|
153
|
+
from warp.types import Volume as Volume
|
|
154
|
+
from warp.types import BvhQuery as BvhQuery
|
|
155
|
+
from warp.types import HashGridQuery as HashGridQuery
|
|
156
|
+
from warp.types import MeshQueryAABB as MeshQueryAABB
|
|
157
|
+
from warp.types import MeshQueryPoint as MeshQueryPoint
|
|
158
|
+
from warp.types import MeshQueryRay as MeshQueryRay
|
|
159
|
+
|
|
160
|
+
from warp.types import matmul as matmul
|
|
161
|
+
from warp.types import adj_matmul as adj_matmul
|
|
162
|
+
from warp.types import batched_matmul as batched_matmul
|
|
163
|
+
from warp.types import adj_batched_matmul as adj_batched_matmul
|
|
43
164
|
|
|
44
|
-
from warp.types import
|
|
45
|
-
from warp.types import
|
|
46
|
-
|
|
47
|
-
from warp.types import
|
|
48
|
-
from warp.types import
|
|
49
|
-
|
|
50
|
-
from warp.types import
|
|
51
|
-
from warp.types import
|
|
52
|
-
|
|
53
|
-
from warp.types import
|
|
54
|
-
|
|
165
|
+
from warp.types import vector as vec
|
|
166
|
+
from warp.types import matrix as mat
|
|
167
|
+
|
|
168
|
+
from warp.types import matrix_from_cols as matrix_from_cols
|
|
169
|
+
from warp.types import matrix_from_rows as matrix_from_rows
|
|
170
|
+
|
|
171
|
+
from warp.types import dtype_from_numpy as dtype_from_numpy
|
|
172
|
+
from warp.types import dtype_to_numpy as dtype_to_numpy
|
|
173
|
+
|
|
174
|
+
from warp.types import from_ipc_handle as from_ipc_handle
|
|
175
|
+
|
|
176
|
+
from warp.context import init as init
|
|
177
|
+
from warp.context import func as func
|
|
178
|
+
from warp.context import func_grad as func_grad
|
|
179
|
+
from warp.context import func_replay as func_replay
|
|
180
|
+
from warp.context import func_native as func_native
|
|
181
|
+
from warp.context import kernel as kernel
|
|
182
|
+
from warp.context import struct as struct
|
|
183
|
+
from warp.context import overload as overload
|
|
184
|
+
|
|
185
|
+
from warp.context import is_cpu_available as is_cpu_available
|
|
186
|
+
from warp.context import is_cuda_available as is_cuda_available
|
|
187
|
+
from warp.context import is_device_available as is_device_available
|
|
188
|
+
from warp.context import get_devices as get_devices
|
|
189
|
+
from warp.context import get_preferred_device as get_preferred_device
|
|
190
|
+
from warp.context import get_cuda_devices as get_cuda_devices
|
|
191
|
+
from warp.context import get_cuda_device_count as get_cuda_device_count
|
|
192
|
+
from warp.context import get_cuda_device as get_cuda_device
|
|
193
|
+
from warp.context import map_cuda_device as map_cuda_device
|
|
194
|
+
from warp.context import unmap_cuda_device as unmap_cuda_device
|
|
195
|
+
from warp.context import get_device as get_device
|
|
196
|
+
from warp.context import set_device as set_device
|
|
197
|
+
from warp.context import synchronize_device as synchronize_device
|
|
198
|
+
|
|
199
|
+
from warp.context import zeros as zeros
|
|
200
|
+
from warp.context import zeros_like as zeros_like
|
|
201
|
+
from warp.context import ones as ones
|
|
202
|
+
from warp.context import ones_like as ones_like
|
|
203
|
+
from warp.context import full as full
|
|
204
|
+
from warp.context import full_like as full_like
|
|
205
|
+
from warp.context import clone as clone
|
|
206
|
+
from warp.context import empty as empty
|
|
207
|
+
from warp.context import empty_like as empty_like
|
|
208
|
+
from warp.context import copy as copy
|
|
209
|
+
from warp.context import from_numpy as from_numpy
|
|
210
|
+
|
|
211
|
+
from warp.context import launch as launch
|
|
212
|
+
from warp.context import launch_tiled as launch_tiled
|
|
213
|
+
from warp.context import synchronize as synchronize
|
|
214
|
+
from warp.context import compile_aot_module as compile_aot_module
|
|
215
|
+
from warp.context import force_load as force_load
|
|
216
|
+
from warp.context import load_module as load_module
|
|
217
|
+
from warp.context import load_aot_module as load_aot_module
|
|
218
|
+
from warp.context import event_from_ipc_handle as event_from_ipc_handle
|
|
219
|
+
|
|
220
|
+
from warp.context import set_module_options as set_module_options
|
|
221
|
+
from warp.context import get_module_options as get_module_options
|
|
222
|
+
from warp.context import get_module as get_module
|
|
223
|
+
|
|
224
|
+
from warp.context import capture_begin as capture_begin
|
|
225
|
+
from warp.context import capture_end as capture_end
|
|
226
|
+
from warp.context import capture_launch as capture_launch
|
|
227
|
+
from warp.context import capture_if as capture_if
|
|
228
|
+
from warp.context import capture_while as capture_while
|
|
229
|
+
from warp.context import capture_debug_dot_print as capture_debug_dot_print
|
|
230
|
+
|
|
231
|
+
from warp.context import Kernel as Kernel
|
|
232
|
+
from warp.context import Function as Function
|
|
233
|
+
from warp.context import Launch as Launch
|
|
234
|
+
|
|
235
|
+
from warp.context import Stream as Stream
|
|
236
|
+
from warp.context import get_stream as get_stream
|
|
237
|
+
from warp.context import set_stream as set_stream
|
|
238
|
+
from warp.context import wait_stream as wait_stream
|
|
239
|
+
from warp.context import synchronize_stream as synchronize_stream
|
|
240
|
+
|
|
241
|
+
from warp.context import Event as Event
|
|
242
|
+
from warp.context import record_event as record_event
|
|
243
|
+
from warp.context import wait_event as wait_event
|
|
244
|
+
from warp.context import synchronize_event as synchronize_event
|
|
245
|
+
from warp.context import get_event_elapsed_time as get_event_elapsed_time
|
|
246
|
+
|
|
247
|
+
from warp.context import RegisteredGLBuffer as RegisteredGLBuffer
|
|
248
|
+
|
|
249
|
+
from warp.context import is_mempool_supported as is_mempool_supported
|
|
250
|
+
from warp.context import is_mempool_enabled as is_mempool_enabled
|
|
251
|
+
from warp.context import set_mempool_enabled as set_mempool_enabled
|
|
252
|
+
|
|
253
|
+
from warp.context import set_mempool_release_threshold as set_mempool_release_threshold
|
|
254
|
+
from warp.context import get_mempool_release_threshold as get_mempool_release_threshold
|
|
255
|
+
from warp.context import get_mempool_used_mem_current as get_mempool_used_mem_current
|
|
256
|
+
from warp.context import get_mempool_used_mem_high as get_mempool_used_mem_high
|
|
257
|
+
|
|
258
|
+
from warp.context import is_mempool_access_supported as is_mempool_access_supported
|
|
259
|
+
from warp.context import is_mempool_access_enabled as is_mempool_access_enabled
|
|
260
|
+
from warp.context import set_mempool_access_enabled as set_mempool_access_enabled
|
|
261
|
+
|
|
262
|
+
from warp.context import is_peer_access_supported as is_peer_access_supported
|
|
263
|
+
from warp.context import is_peer_access_enabled as is_peer_access_enabled
|
|
264
|
+
from warp.context import set_peer_access_enabled as set_peer_access_enabled
|
|
265
|
+
|
|
266
|
+
from warp.tape import Tape as Tape
|
|
267
|
+
|
|
268
|
+
from warp.utils import ScopedTimer as ScopedTimer
|
|
269
|
+
from warp.utils import ScopedDevice as ScopedDevice
|
|
270
|
+
from warp.utils import ScopedStream as ScopedStream
|
|
271
|
+
from warp.utils import ScopedMempool as ScopedMempool
|
|
272
|
+
from warp.utils import ScopedMempoolAccess as ScopedMempoolAccess
|
|
273
|
+
from warp.utils import ScopedPeerAccess as ScopedPeerAccess
|
|
274
|
+
from warp.utils import ScopedCapture as ScopedCapture
|
|
275
|
+
|
|
276
|
+
from warp.utils import transform_expand as transform_expand
|
|
277
|
+
from warp.utils import quat_between_vectors as quat_between_vectors
|
|
278
|
+
|
|
279
|
+
from warp.utils import TimingResult as TimingResult
|
|
280
|
+
from warp.utils import timing_begin as timing_begin
|
|
281
|
+
from warp.utils import timing_end as timing_end
|
|
282
|
+
from warp.utils import timing_print as timing_print
|
|
283
|
+
|
|
284
|
+
from warp.utils import TIMING_KERNEL as TIMING_KERNEL
|
|
285
|
+
from warp.utils import TIMING_KERNEL_BUILTIN as TIMING_KERNEL_BUILTIN
|
|
286
|
+
from warp.utils import TIMING_MEMCPY as TIMING_MEMCPY
|
|
287
|
+
from warp.utils import TIMING_MEMSET as TIMING_MEMSET
|
|
288
|
+
from warp.utils import TIMING_GRAPH as TIMING_GRAPH
|
|
289
|
+
from warp.utils import TIMING_ALL as TIMING_ALL
|
|
290
|
+
|
|
291
|
+
from warp.utils import map as map
|
|
292
|
+
|
|
293
|
+
from warp.marching_cubes import MarchingCubes as MarchingCubes
|
|
294
|
+
|
|
295
|
+
from warp.torch import from_torch as from_torch
|
|
296
|
+
from warp.torch import to_torch as to_torch
|
|
297
|
+
from warp.torch import dtype_from_torch as dtype_from_torch
|
|
298
|
+
from warp.torch import dtype_to_torch as dtype_to_torch
|
|
299
|
+
from warp.torch import device_from_torch as device_from_torch
|
|
300
|
+
from warp.torch import device_to_torch as device_to_torch
|
|
301
|
+
from warp.torch import stream_from_torch as stream_from_torch
|
|
302
|
+
from warp.torch import stream_to_torch as stream_to_torch
|
|
303
|
+
|
|
304
|
+
from warp.jax import from_jax as from_jax
|
|
305
|
+
from warp.jax import to_jax as to_jax
|
|
306
|
+
from warp.jax import dtype_from_jax as dtype_from_jax
|
|
307
|
+
from warp.jax import dtype_to_jax as dtype_to_jax
|
|
308
|
+
from warp.jax import device_from_jax as device_from_jax
|
|
309
|
+
from warp.jax import device_to_jax as device_to_jax
|
|
310
|
+
|
|
311
|
+
from warp.dlpack import from_dlpack as from_dlpack
|
|
312
|
+
from warp.dlpack import to_dlpack as to_dlpack
|
|
313
|
+
|
|
314
|
+
from warp.paddle import from_paddle as from_paddle
|
|
315
|
+
from warp.paddle import to_paddle as to_paddle
|
|
316
|
+
from warp.paddle import dtype_from_paddle as dtype_from_paddle
|
|
317
|
+
from warp.paddle import dtype_to_paddle as dtype_to_paddle
|
|
318
|
+
from warp.paddle import device_from_paddle as device_from_paddle
|
|
319
|
+
from warp.paddle import device_to_paddle as device_to_paddle
|
|
320
|
+
from warp.paddle import stream_from_paddle as stream_from_paddle
|
|
321
|
+
|
|
322
|
+
from warp.build import clear_kernel_cache as clear_kernel_cache
|
|
323
|
+
from warp.build import clear_lto_cache as clear_lto_cache
|
|
324
|
+
|
|
325
|
+
from warp.constants import *
|
|
326
|
+
|
|
327
|
+
from . import builtins
|
|
328
|
+
from warp.builtins import static as static
|
|
329
|
+
|
|
330
|
+
from warp.math import *
|
|
331
|
+
|
|
332
|
+
from . import config as config
|
|
333
|
+
|
|
334
|
+
__version__ = config.version
|
|
335
|
+
|
|
336
|
+
class vec2h:
|
|
337
|
+
@over
|
|
338
|
+
def __init__(self) -> None:
|
|
339
|
+
"""Construct a zero-initialized vector."""
|
|
340
|
+
...
|
|
341
|
+
|
|
342
|
+
@over
|
|
343
|
+
def __init__(self, other: vec2h) -> None:
|
|
344
|
+
"""Construct a vector by copy."""
|
|
345
|
+
...
|
|
346
|
+
|
|
347
|
+
@over
|
|
348
|
+
def __init__(self, x: float16, y: float16) -> None:
|
|
349
|
+
"""Construct a vector from its component values."""
|
|
350
|
+
...
|
|
351
|
+
|
|
352
|
+
@over
|
|
353
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
354
|
+
"""Construct a vector from a sequence of values."""
|
|
355
|
+
...
|
|
356
|
+
|
|
357
|
+
@over
|
|
358
|
+
def __init__(self, value: float16) -> None:
|
|
359
|
+
"""Construct a vector filled with a value."""
|
|
360
|
+
...
|
|
361
|
+
|
|
362
|
+
class vec2f:
|
|
363
|
+
@over
|
|
364
|
+
def __init__(self) -> None:
|
|
365
|
+
"""Construct a zero-initialized vector."""
|
|
366
|
+
...
|
|
367
|
+
|
|
368
|
+
@over
|
|
369
|
+
def __init__(self, other: vec2f) -> None:
|
|
370
|
+
"""Construct a vector by copy."""
|
|
371
|
+
...
|
|
372
|
+
|
|
373
|
+
@over
|
|
374
|
+
def __init__(self, x: float32, y: float32) -> None:
|
|
375
|
+
"""Construct a vector from its component values."""
|
|
376
|
+
...
|
|
377
|
+
|
|
378
|
+
@over
|
|
379
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
380
|
+
"""Construct a vector from a sequence of values."""
|
|
381
|
+
...
|
|
382
|
+
|
|
383
|
+
@over
|
|
384
|
+
def __init__(self, value: float32) -> None:
|
|
385
|
+
"""Construct a vector filled with a value."""
|
|
386
|
+
...
|
|
387
|
+
|
|
388
|
+
class vec2d:
|
|
389
|
+
@over
|
|
390
|
+
def __init__(self) -> None:
|
|
391
|
+
"""Construct a zero-initialized vector."""
|
|
392
|
+
...
|
|
393
|
+
|
|
394
|
+
@over
|
|
395
|
+
def __init__(self, other: vec2d) -> None:
|
|
396
|
+
"""Construct a vector by copy."""
|
|
397
|
+
...
|
|
398
|
+
|
|
399
|
+
@over
|
|
400
|
+
def __init__(self, x: float64, y: float64) -> None:
|
|
401
|
+
"""Construct a vector from its component values."""
|
|
402
|
+
...
|
|
403
|
+
|
|
404
|
+
@over
|
|
405
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
406
|
+
"""Construct a vector from a sequence of values."""
|
|
407
|
+
...
|
|
408
|
+
|
|
409
|
+
@over
|
|
410
|
+
def __init__(self, value: float64) -> None:
|
|
411
|
+
"""Construct a vector filled with a value."""
|
|
412
|
+
...
|
|
413
|
+
|
|
414
|
+
class vec2b:
|
|
415
|
+
@over
|
|
416
|
+
def __init__(self) -> None:
|
|
417
|
+
"""Construct a zero-initialized vector."""
|
|
418
|
+
...
|
|
419
|
+
|
|
420
|
+
@over
|
|
421
|
+
def __init__(self, other: vec2b) -> None:
|
|
422
|
+
"""Construct a vector by copy."""
|
|
423
|
+
...
|
|
424
|
+
|
|
425
|
+
@over
|
|
426
|
+
def __init__(self, x: int8, y: int8) -> None:
|
|
427
|
+
"""Construct a vector from its component values."""
|
|
428
|
+
...
|
|
429
|
+
|
|
430
|
+
@over
|
|
431
|
+
def __init__(self, args: Sequence[int8]) -> None:
|
|
432
|
+
"""Construct a vector from a sequence of values."""
|
|
433
|
+
...
|
|
434
|
+
|
|
435
|
+
@over
|
|
436
|
+
def __init__(self, value: int8) -> None:
|
|
437
|
+
"""Construct a vector filled with a value."""
|
|
438
|
+
...
|
|
439
|
+
|
|
440
|
+
class vec2ub:
|
|
441
|
+
@over
|
|
442
|
+
def __init__(self) -> None:
|
|
443
|
+
"""Construct a zero-initialized vector."""
|
|
444
|
+
...
|
|
445
|
+
|
|
446
|
+
@over
|
|
447
|
+
def __init__(self, other: vec2ub) -> None:
|
|
448
|
+
"""Construct a vector by copy."""
|
|
449
|
+
...
|
|
450
|
+
|
|
451
|
+
@over
|
|
452
|
+
def __init__(self, x: uint8, y: uint8) -> None:
|
|
453
|
+
"""Construct a vector from its component values."""
|
|
454
|
+
...
|
|
455
|
+
|
|
456
|
+
@over
|
|
457
|
+
def __init__(self, args: Sequence[uint8]) -> None:
|
|
458
|
+
"""Construct a vector from a sequence of values."""
|
|
459
|
+
...
|
|
460
|
+
|
|
461
|
+
@over
|
|
462
|
+
def __init__(self, value: uint8) -> None:
|
|
463
|
+
"""Construct a vector filled with a value."""
|
|
464
|
+
...
|
|
465
|
+
|
|
466
|
+
class vec2s:
|
|
467
|
+
@over
|
|
468
|
+
def __init__(self) -> None:
|
|
469
|
+
"""Construct a zero-initialized vector."""
|
|
470
|
+
...
|
|
471
|
+
|
|
472
|
+
@over
|
|
473
|
+
def __init__(self, other: vec2s) -> None:
|
|
474
|
+
"""Construct a vector by copy."""
|
|
475
|
+
...
|
|
476
|
+
|
|
477
|
+
@over
|
|
478
|
+
def __init__(self, x: int16, y: int16) -> None:
|
|
479
|
+
"""Construct a vector from its component values."""
|
|
480
|
+
...
|
|
481
|
+
|
|
482
|
+
@over
|
|
483
|
+
def __init__(self, args: Sequence[int16]) -> None:
|
|
484
|
+
"""Construct a vector from a sequence of values."""
|
|
485
|
+
...
|
|
486
|
+
|
|
487
|
+
@over
|
|
488
|
+
def __init__(self, value: int16) -> None:
|
|
489
|
+
"""Construct a vector filled with a value."""
|
|
490
|
+
...
|
|
491
|
+
|
|
492
|
+
class vec2us:
|
|
493
|
+
@over
|
|
494
|
+
def __init__(self) -> None:
|
|
495
|
+
"""Construct a zero-initialized vector."""
|
|
496
|
+
...
|
|
497
|
+
|
|
498
|
+
@over
|
|
499
|
+
def __init__(self, other: vec2us) -> None:
|
|
500
|
+
"""Construct a vector by copy."""
|
|
501
|
+
...
|
|
502
|
+
|
|
503
|
+
@over
|
|
504
|
+
def __init__(self, x: uint16, y: uint16) -> None:
|
|
505
|
+
"""Construct a vector from its component values."""
|
|
506
|
+
...
|
|
507
|
+
|
|
508
|
+
@over
|
|
509
|
+
def __init__(self, args: Sequence[uint16]) -> None:
|
|
510
|
+
"""Construct a vector from a sequence of values."""
|
|
511
|
+
...
|
|
512
|
+
|
|
513
|
+
@over
|
|
514
|
+
def __init__(self, value: uint16) -> None:
|
|
515
|
+
"""Construct a vector filled with a value."""
|
|
516
|
+
...
|
|
517
|
+
|
|
518
|
+
class vec2i:
|
|
519
|
+
@over
|
|
520
|
+
def __init__(self) -> None:
|
|
521
|
+
"""Construct a zero-initialized vector."""
|
|
522
|
+
...
|
|
523
|
+
|
|
524
|
+
@over
|
|
525
|
+
def __init__(self, other: vec2i) -> None:
|
|
526
|
+
"""Construct a vector by copy."""
|
|
527
|
+
...
|
|
528
|
+
|
|
529
|
+
@over
|
|
530
|
+
def __init__(self, x: int32, y: int32) -> None:
|
|
531
|
+
"""Construct a vector from its component values."""
|
|
532
|
+
...
|
|
533
|
+
|
|
534
|
+
@over
|
|
535
|
+
def __init__(self, args: Sequence[int32]) -> None:
|
|
536
|
+
"""Construct a vector from a sequence of values."""
|
|
537
|
+
...
|
|
538
|
+
|
|
539
|
+
@over
|
|
540
|
+
def __init__(self, value: int32) -> None:
|
|
541
|
+
"""Construct a vector filled with a value."""
|
|
542
|
+
...
|
|
543
|
+
|
|
544
|
+
class vec2ui:
|
|
545
|
+
@over
|
|
546
|
+
def __init__(self) -> None:
|
|
547
|
+
"""Construct a zero-initialized vector."""
|
|
548
|
+
...
|
|
549
|
+
|
|
550
|
+
@over
|
|
551
|
+
def __init__(self, other: vec2ui) -> None:
|
|
552
|
+
"""Construct a vector by copy."""
|
|
553
|
+
...
|
|
554
|
+
|
|
555
|
+
@over
|
|
556
|
+
def __init__(self, x: uint32, y: uint32) -> None:
|
|
557
|
+
"""Construct a vector from its component values."""
|
|
558
|
+
...
|
|
559
|
+
|
|
560
|
+
@over
|
|
561
|
+
def __init__(self, args: Sequence[uint32]) -> None:
|
|
562
|
+
"""Construct a vector from a sequence of values."""
|
|
563
|
+
...
|
|
564
|
+
|
|
565
|
+
@over
|
|
566
|
+
def __init__(self, value: uint32) -> None:
|
|
567
|
+
"""Construct a vector filled with a value."""
|
|
568
|
+
...
|
|
569
|
+
|
|
570
|
+
class vec2l:
|
|
571
|
+
@over
|
|
572
|
+
def __init__(self) -> None:
|
|
573
|
+
"""Construct a zero-initialized vector."""
|
|
574
|
+
...
|
|
575
|
+
|
|
576
|
+
@over
|
|
577
|
+
def __init__(self, other: vec2l) -> None:
|
|
578
|
+
"""Construct a vector by copy."""
|
|
579
|
+
...
|
|
580
|
+
|
|
581
|
+
@over
|
|
582
|
+
def __init__(self, x: int64, y: int64) -> None:
|
|
583
|
+
"""Construct a vector from its component values."""
|
|
584
|
+
...
|
|
585
|
+
|
|
586
|
+
@over
|
|
587
|
+
def __init__(self, args: Sequence[int64]) -> None:
|
|
588
|
+
"""Construct a vector from a sequence of values."""
|
|
589
|
+
...
|
|
590
|
+
|
|
591
|
+
@over
|
|
592
|
+
def __init__(self, value: int64) -> None:
|
|
593
|
+
"""Construct a vector filled with a value."""
|
|
594
|
+
...
|
|
595
|
+
|
|
596
|
+
class vec2ul:
|
|
597
|
+
@over
|
|
598
|
+
def __init__(self) -> None:
|
|
599
|
+
"""Construct a zero-initialized vector."""
|
|
600
|
+
...
|
|
601
|
+
|
|
602
|
+
@over
|
|
603
|
+
def __init__(self, other: vec2ul) -> None:
|
|
604
|
+
"""Construct a vector by copy."""
|
|
605
|
+
...
|
|
606
|
+
|
|
607
|
+
@over
|
|
608
|
+
def __init__(self, x: uint64, y: uint64) -> None:
|
|
609
|
+
"""Construct a vector from its component values."""
|
|
610
|
+
...
|
|
611
|
+
|
|
612
|
+
@over
|
|
613
|
+
def __init__(self, args: Sequence[uint64]) -> None:
|
|
614
|
+
"""Construct a vector from a sequence of values."""
|
|
615
|
+
...
|
|
616
|
+
|
|
617
|
+
@over
|
|
618
|
+
def __init__(self, value: uint64) -> None:
|
|
619
|
+
"""Construct a vector filled with a value."""
|
|
620
|
+
...
|
|
621
|
+
|
|
622
|
+
vec2 = vec2f
|
|
623
|
+
|
|
624
|
+
class vec3h:
|
|
625
|
+
@over
|
|
626
|
+
def __init__(self) -> None:
|
|
627
|
+
"""Construct a zero-initialized vector."""
|
|
628
|
+
...
|
|
629
|
+
|
|
630
|
+
@over
|
|
631
|
+
def __init__(self, other: vec3h) -> None:
|
|
632
|
+
"""Construct a vector by copy."""
|
|
633
|
+
...
|
|
634
|
+
|
|
635
|
+
@over
|
|
636
|
+
def __init__(self, x: float16, y: float16, z: float16) -> None:
|
|
637
|
+
"""Construct a vector from its component values."""
|
|
638
|
+
...
|
|
639
|
+
|
|
640
|
+
@over
|
|
641
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
642
|
+
"""Construct a vector from a sequence of values."""
|
|
643
|
+
...
|
|
644
|
+
|
|
645
|
+
@over
|
|
646
|
+
def __init__(self, value: float16) -> None:
|
|
647
|
+
"""Construct a vector filled with a value."""
|
|
648
|
+
...
|
|
649
|
+
|
|
650
|
+
class vec3f:
|
|
651
|
+
@over
|
|
652
|
+
def __init__(self) -> None:
|
|
653
|
+
"""Construct a zero-initialized vector."""
|
|
654
|
+
...
|
|
655
|
+
|
|
656
|
+
@over
|
|
657
|
+
def __init__(self, other: vec3f) -> None:
|
|
658
|
+
"""Construct a vector by copy."""
|
|
659
|
+
...
|
|
660
|
+
|
|
661
|
+
@over
|
|
662
|
+
def __init__(self, x: float32, y: float32, z: float32) -> None:
|
|
663
|
+
"""Construct a vector from its component values."""
|
|
664
|
+
...
|
|
665
|
+
|
|
666
|
+
@over
|
|
667
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
668
|
+
"""Construct a vector from a sequence of values."""
|
|
669
|
+
...
|
|
670
|
+
|
|
671
|
+
@over
|
|
672
|
+
def __init__(self, value: float32) -> None:
|
|
673
|
+
"""Construct a vector filled with a value."""
|
|
674
|
+
...
|
|
675
|
+
|
|
676
|
+
class vec3d:
|
|
677
|
+
@over
|
|
678
|
+
def __init__(self) -> None:
|
|
679
|
+
"""Construct a zero-initialized vector."""
|
|
680
|
+
...
|
|
681
|
+
|
|
682
|
+
@over
|
|
683
|
+
def __init__(self, other: vec3d) -> None:
|
|
684
|
+
"""Construct a vector by copy."""
|
|
685
|
+
...
|
|
686
|
+
|
|
687
|
+
@over
|
|
688
|
+
def __init__(self, x: float64, y: float64, z: float64) -> None:
|
|
689
|
+
"""Construct a vector from its component values."""
|
|
690
|
+
...
|
|
691
|
+
|
|
692
|
+
@over
|
|
693
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
694
|
+
"""Construct a vector from a sequence of values."""
|
|
695
|
+
...
|
|
696
|
+
|
|
697
|
+
@over
|
|
698
|
+
def __init__(self, value: float64) -> None:
|
|
699
|
+
"""Construct a vector filled with a value."""
|
|
700
|
+
...
|
|
701
|
+
|
|
702
|
+
class vec3b:
|
|
703
|
+
@over
|
|
704
|
+
def __init__(self) -> None:
|
|
705
|
+
"""Construct a zero-initialized vector."""
|
|
706
|
+
...
|
|
707
|
+
|
|
708
|
+
@over
|
|
709
|
+
def __init__(self, other: vec3b) -> None:
|
|
710
|
+
"""Construct a vector by copy."""
|
|
711
|
+
...
|
|
712
|
+
|
|
713
|
+
@over
|
|
714
|
+
def __init__(self, x: int8, y: int8, z: int8) -> None:
|
|
715
|
+
"""Construct a vector from its component values."""
|
|
716
|
+
...
|
|
717
|
+
|
|
718
|
+
@over
|
|
719
|
+
def __init__(self, args: Sequence[int8]) -> None:
|
|
720
|
+
"""Construct a vector from a sequence of values."""
|
|
721
|
+
...
|
|
722
|
+
|
|
723
|
+
@over
|
|
724
|
+
def __init__(self, value: int8) -> None:
|
|
725
|
+
"""Construct a vector filled with a value."""
|
|
726
|
+
...
|
|
727
|
+
|
|
728
|
+
class vec3ub:
|
|
729
|
+
@over
|
|
730
|
+
def __init__(self) -> None:
|
|
731
|
+
"""Construct a zero-initialized vector."""
|
|
732
|
+
...
|
|
733
|
+
|
|
734
|
+
@over
|
|
735
|
+
def __init__(self, other: vec3ub) -> None:
|
|
736
|
+
"""Construct a vector by copy."""
|
|
737
|
+
...
|
|
738
|
+
|
|
739
|
+
@over
|
|
740
|
+
def __init__(self, x: uint8, y: uint8, z: uint8) -> None:
|
|
741
|
+
"""Construct a vector from its component values."""
|
|
742
|
+
...
|
|
743
|
+
|
|
744
|
+
@over
|
|
745
|
+
def __init__(self, args: Sequence[uint8]) -> None:
|
|
746
|
+
"""Construct a vector from a sequence of values."""
|
|
747
|
+
...
|
|
748
|
+
|
|
749
|
+
@over
|
|
750
|
+
def __init__(self, value: uint8) -> None:
|
|
751
|
+
"""Construct a vector filled with a value."""
|
|
752
|
+
...
|
|
753
|
+
|
|
754
|
+
class vec3s:
|
|
755
|
+
@over
|
|
756
|
+
def __init__(self) -> None:
|
|
757
|
+
"""Construct a zero-initialized vector."""
|
|
758
|
+
...
|
|
759
|
+
|
|
760
|
+
@over
|
|
761
|
+
def __init__(self, other: vec3s) -> None:
|
|
762
|
+
"""Construct a vector by copy."""
|
|
763
|
+
...
|
|
764
|
+
|
|
765
|
+
@over
|
|
766
|
+
def __init__(self, x: int16, y: int16, z: int16) -> None:
|
|
767
|
+
"""Construct a vector from its component values."""
|
|
768
|
+
...
|
|
769
|
+
|
|
770
|
+
@over
|
|
771
|
+
def __init__(self, args: Sequence[int16]) -> None:
|
|
772
|
+
"""Construct a vector from a sequence of values."""
|
|
773
|
+
...
|
|
774
|
+
|
|
775
|
+
@over
|
|
776
|
+
def __init__(self, value: int16) -> None:
|
|
777
|
+
"""Construct a vector filled with a value."""
|
|
778
|
+
...
|
|
779
|
+
|
|
780
|
+
class vec3us:
|
|
781
|
+
@over
|
|
782
|
+
def __init__(self) -> None:
|
|
783
|
+
"""Construct a zero-initialized vector."""
|
|
784
|
+
...
|
|
785
|
+
|
|
786
|
+
@over
|
|
787
|
+
def __init__(self, other: vec3us) -> None:
|
|
788
|
+
"""Construct a vector by copy."""
|
|
789
|
+
...
|
|
790
|
+
|
|
791
|
+
@over
|
|
792
|
+
def __init__(self, x: uint16, y: uint16, z: uint16) -> None:
|
|
793
|
+
"""Construct a vector from its component values."""
|
|
794
|
+
...
|
|
795
|
+
|
|
796
|
+
@over
|
|
797
|
+
def __init__(self, args: Sequence[uint16]) -> None:
|
|
798
|
+
"""Construct a vector from a sequence of values."""
|
|
799
|
+
...
|
|
800
|
+
|
|
801
|
+
@over
|
|
802
|
+
def __init__(self, value: uint16) -> None:
|
|
803
|
+
"""Construct a vector filled with a value."""
|
|
804
|
+
...
|
|
805
|
+
|
|
806
|
+
class vec3i:
|
|
807
|
+
@over
|
|
808
|
+
def __init__(self) -> None:
|
|
809
|
+
"""Construct a zero-initialized vector."""
|
|
810
|
+
...
|
|
811
|
+
|
|
812
|
+
@over
|
|
813
|
+
def __init__(self, other: vec3i) -> None:
|
|
814
|
+
"""Construct a vector by copy."""
|
|
815
|
+
...
|
|
816
|
+
|
|
817
|
+
@over
|
|
818
|
+
def __init__(self, x: int32, y: int32, z: int32) -> None:
|
|
819
|
+
"""Construct a vector from its component values."""
|
|
820
|
+
...
|
|
821
|
+
|
|
822
|
+
@over
|
|
823
|
+
def __init__(self, args: Sequence[int32]) -> None:
|
|
824
|
+
"""Construct a vector from a sequence of values."""
|
|
825
|
+
...
|
|
826
|
+
|
|
827
|
+
@over
|
|
828
|
+
def __init__(self, value: int32) -> None:
|
|
829
|
+
"""Construct a vector filled with a value."""
|
|
830
|
+
...
|
|
831
|
+
|
|
832
|
+
class vec3ui:
|
|
833
|
+
@over
|
|
834
|
+
def __init__(self) -> None:
|
|
835
|
+
"""Construct a zero-initialized vector."""
|
|
836
|
+
...
|
|
837
|
+
|
|
838
|
+
@over
|
|
839
|
+
def __init__(self, other: vec3ui) -> None:
|
|
840
|
+
"""Construct a vector by copy."""
|
|
841
|
+
...
|
|
842
|
+
|
|
843
|
+
@over
|
|
844
|
+
def __init__(self, x: uint32, y: uint32, z: uint32) -> None:
|
|
845
|
+
"""Construct a vector from its component values."""
|
|
846
|
+
...
|
|
847
|
+
|
|
848
|
+
@over
|
|
849
|
+
def __init__(self, args: Sequence[uint32]) -> None:
|
|
850
|
+
"""Construct a vector from a sequence of values."""
|
|
851
|
+
...
|
|
852
|
+
|
|
853
|
+
@over
|
|
854
|
+
def __init__(self, value: uint32) -> None:
|
|
855
|
+
"""Construct a vector filled with a value."""
|
|
856
|
+
...
|
|
857
|
+
|
|
858
|
+
class vec3l:
|
|
859
|
+
@over
|
|
860
|
+
def __init__(self) -> None:
|
|
861
|
+
"""Construct a zero-initialized vector."""
|
|
862
|
+
...
|
|
863
|
+
|
|
864
|
+
@over
|
|
865
|
+
def __init__(self, other: vec3l) -> None:
|
|
866
|
+
"""Construct a vector by copy."""
|
|
867
|
+
...
|
|
868
|
+
|
|
869
|
+
@over
|
|
870
|
+
def __init__(self, x: int64, y: int64, z: int64) -> None:
|
|
871
|
+
"""Construct a vector from its component values."""
|
|
872
|
+
...
|
|
873
|
+
|
|
874
|
+
@over
|
|
875
|
+
def __init__(self, args: Sequence[int64]) -> None:
|
|
876
|
+
"""Construct a vector from a sequence of values."""
|
|
877
|
+
...
|
|
878
|
+
|
|
879
|
+
@over
|
|
880
|
+
def __init__(self, value: int64) -> None:
|
|
881
|
+
"""Construct a vector filled with a value."""
|
|
882
|
+
...
|
|
883
|
+
|
|
884
|
+
class vec3ul:
|
|
885
|
+
@over
|
|
886
|
+
def __init__(self) -> None:
|
|
887
|
+
"""Construct a zero-initialized vector."""
|
|
888
|
+
...
|
|
889
|
+
|
|
890
|
+
@over
|
|
891
|
+
def __init__(self, other: vec3ul) -> None:
|
|
892
|
+
"""Construct a vector by copy."""
|
|
893
|
+
...
|
|
894
|
+
|
|
895
|
+
@over
|
|
896
|
+
def __init__(self, x: uint64, y: uint64, z: uint64) -> None:
|
|
897
|
+
"""Construct a vector from its component values."""
|
|
898
|
+
...
|
|
899
|
+
|
|
900
|
+
@over
|
|
901
|
+
def __init__(self, args: Sequence[uint64]) -> None:
|
|
902
|
+
"""Construct a vector from a sequence of values."""
|
|
903
|
+
...
|
|
904
|
+
|
|
905
|
+
@over
|
|
906
|
+
def __init__(self, value: uint64) -> None:
|
|
907
|
+
"""Construct a vector filled with a value."""
|
|
908
|
+
...
|
|
909
|
+
|
|
910
|
+
vec3 = vec3f
|
|
911
|
+
|
|
912
|
+
class vec4h:
|
|
913
|
+
@over
|
|
914
|
+
def __init__(self) -> None:
|
|
915
|
+
"""Construct a zero-initialized vector."""
|
|
916
|
+
...
|
|
917
|
+
|
|
918
|
+
@over
|
|
919
|
+
def __init__(self, other: vec4h) -> None:
|
|
920
|
+
"""Construct a vector by copy."""
|
|
921
|
+
...
|
|
922
|
+
|
|
923
|
+
@over
|
|
924
|
+
def __init__(self, x: float16, y: float16, z: float16, w: float16) -> None:
|
|
925
|
+
"""Construct a vector from its component values."""
|
|
926
|
+
...
|
|
927
|
+
|
|
928
|
+
@over
|
|
929
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
930
|
+
"""Construct a vector from a sequence of values."""
|
|
931
|
+
...
|
|
932
|
+
|
|
933
|
+
@over
|
|
934
|
+
def __init__(self, value: float16) -> None:
|
|
935
|
+
"""Construct a vector filled with a value."""
|
|
936
|
+
...
|
|
937
|
+
|
|
938
|
+
class vec4f:
|
|
939
|
+
@over
|
|
940
|
+
def __init__(self) -> None:
|
|
941
|
+
"""Construct a zero-initialized vector."""
|
|
942
|
+
...
|
|
943
|
+
|
|
944
|
+
@over
|
|
945
|
+
def __init__(self, other: vec4f) -> None:
|
|
946
|
+
"""Construct a vector by copy."""
|
|
947
|
+
...
|
|
948
|
+
|
|
949
|
+
@over
|
|
950
|
+
def __init__(self, x: float32, y: float32, z: float32, w: float32) -> None:
|
|
951
|
+
"""Construct a vector from its component values."""
|
|
952
|
+
...
|
|
953
|
+
|
|
954
|
+
@over
|
|
955
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
956
|
+
"""Construct a vector from a sequence of values."""
|
|
957
|
+
...
|
|
958
|
+
|
|
959
|
+
@over
|
|
960
|
+
def __init__(self, value: float32) -> None:
|
|
961
|
+
"""Construct a vector filled with a value."""
|
|
962
|
+
...
|
|
963
|
+
|
|
964
|
+
class vec4d:
|
|
965
|
+
@over
|
|
966
|
+
def __init__(self) -> None:
|
|
967
|
+
"""Construct a zero-initialized vector."""
|
|
968
|
+
...
|
|
969
|
+
|
|
970
|
+
@over
|
|
971
|
+
def __init__(self, other: vec4d) -> None:
|
|
972
|
+
"""Construct a vector by copy."""
|
|
973
|
+
...
|
|
974
|
+
|
|
975
|
+
@over
|
|
976
|
+
def __init__(self, x: float64, y: float64, z: float64, w: float64) -> None:
|
|
977
|
+
"""Construct a vector from its component values."""
|
|
978
|
+
...
|
|
979
|
+
|
|
980
|
+
@over
|
|
981
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
982
|
+
"""Construct a vector from a sequence of values."""
|
|
983
|
+
...
|
|
984
|
+
|
|
985
|
+
@over
|
|
986
|
+
def __init__(self, value: float64) -> None:
|
|
987
|
+
"""Construct a vector filled with a value."""
|
|
988
|
+
...
|
|
989
|
+
|
|
990
|
+
class vec4b:
|
|
991
|
+
@over
|
|
992
|
+
def __init__(self) -> None:
|
|
993
|
+
"""Construct a zero-initialized vector."""
|
|
994
|
+
...
|
|
995
|
+
|
|
996
|
+
@over
|
|
997
|
+
def __init__(self, other: vec4b) -> None:
|
|
998
|
+
"""Construct a vector by copy."""
|
|
999
|
+
...
|
|
1000
|
+
|
|
1001
|
+
@over
|
|
1002
|
+
def __init__(self, x: int8, y: int8, z: int8, w: int8) -> None:
|
|
1003
|
+
"""Construct a vector from its component values."""
|
|
1004
|
+
...
|
|
1005
|
+
|
|
1006
|
+
@over
|
|
1007
|
+
def __init__(self, args: Sequence[int8]) -> None:
|
|
1008
|
+
"""Construct a vector from a sequence of values."""
|
|
1009
|
+
...
|
|
1010
|
+
|
|
1011
|
+
@over
|
|
1012
|
+
def __init__(self, value: int8) -> None:
|
|
1013
|
+
"""Construct a vector filled with a value."""
|
|
1014
|
+
...
|
|
1015
|
+
|
|
1016
|
+
class vec4ub:
|
|
1017
|
+
@over
|
|
1018
|
+
def __init__(self) -> None:
|
|
1019
|
+
"""Construct a zero-initialized vector."""
|
|
1020
|
+
...
|
|
1021
|
+
|
|
1022
|
+
@over
|
|
1023
|
+
def __init__(self, other: vec4ub) -> None:
|
|
1024
|
+
"""Construct a vector by copy."""
|
|
1025
|
+
...
|
|
1026
|
+
|
|
1027
|
+
@over
|
|
1028
|
+
def __init__(self, x: uint8, y: uint8, z: uint8, w: uint8) -> None:
|
|
1029
|
+
"""Construct a vector from its component values."""
|
|
1030
|
+
...
|
|
1031
|
+
|
|
1032
|
+
@over
|
|
1033
|
+
def __init__(self, args: Sequence[uint8]) -> None:
|
|
1034
|
+
"""Construct a vector from a sequence of values."""
|
|
1035
|
+
...
|
|
1036
|
+
|
|
1037
|
+
@over
|
|
1038
|
+
def __init__(self, value: uint8) -> None:
|
|
1039
|
+
"""Construct a vector filled with a value."""
|
|
1040
|
+
...
|
|
1041
|
+
|
|
1042
|
+
class vec4s:
|
|
1043
|
+
@over
|
|
1044
|
+
def __init__(self) -> None:
|
|
1045
|
+
"""Construct a zero-initialized vector."""
|
|
1046
|
+
...
|
|
1047
|
+
|
|
1048
|
+
@over
|
|
1049
|
+
def __init__(self, other: vec4s) -> None:
|
|
1050
|
+
"""Construct a vector by copy."""
|
|
1051
|
+
...
|
|
1052
|
+
|
|
1053
|
+
@over
|
|
1054
|
+
def __init__(self, x: int16, y: int16, z: int16, w: int16) -> None:
|
|
1055
|
+
"""Construct a vector from its component values."""
|
|
1056
|
+
...
|
|
1057
|
+
|
|
1058
|
+
@over
|
|
1059
|
+
def __init__(self, args: Sequence[int16]) -> None:
|
|
1060
|
+
"""Construct a vector from a sequence of values."""
|
|
1061
|
+
...
|
|
1062
|
+
|
|
1063
|
+
@over
|
|
1064
|
+
def __init__(self, value: int16) -> None:
|
|
1065
|
+
"""Construct a vector filled with a value."""
|
|
1066
|
+
...
|
|
1067
|
+
|
|
1068
|
+
class vec4us:
|
|
1069
|
+
@over
|
|
1070
|
+
def __init__(self) -> None:
|
|
1071
|
+
"""Construct a zero-initialized vector."""
|
|
1072
|
+
...
|
|
1073
|
+
|
|
1074
|
+
@over
|
|
1075
|
+
def __init__(self, other: vec4us) -> None:
|
|
1076
|
+
"""Construct a vector by copy."""
|
|
1077
|
+
...
|
|
1078
|
+
|
|
1079
|
+
@over
|
|
1080
|
+
def __init__(self, x: uint16, y: uint16, z: uint16, w: uint16) -> None:
|
|
1081
|
+
"""Construct a vector from its component values."""
|
|
1082
|
+
...
|
|
1083
|
+
|
|
1084
|
+
@over
|
|
1085
|
+
def __init__(self, args: Sequence[uint16]) -> None:
|
|
1086
|
+
"""Construct a vector from a sequence of values."""
|
|
1087
|
+
...
|
|
1088
|
+
|
|
1089
|
+
@over
|
|
1090
|
+
def __init__(self, value: uint16) -> None:
|
|
1091
|
+
"""Construct a vector filled with a value."""
|
|
1092
|
+
...
|
|
1093
|
+
|
|
1094
|
+
class vec4i:
|
|
1095
|
+
@over
|
|
1096
|
+
def __init__(self) -> None:
|
|
1097
|
+
"""Construct a zero-initialized vector."""
|
|
1098
|
+
...
|
|
1099
|
+
|
|
1100
|
+
@over
|
|
1101
|
+
def __init__(self, other: vec4i) -> None:
|
|
1102
|
+
"""Construct a vector by copy."""
|
|
1103
|
+
...
|
|
1104
|
+
|
|
1105
|
+
@over
|
|
1106
|
+
def __init__(self, x: int32, y: int32, z: int32, w: int32) -> None:
|
|
1107
|
+
"""Construct a vector from its component values."""
|
|
1108
|
+
...
|
|
1109
|
+
|
|
1110
|
+
@over
|
|
1111
|
+
def __init__(self, args: Sequence[int32]) -> None:
|
|
1112
|
+
"""Construct a vector from a sequence of values."""
|
|
1113
|
+
...
|
|
1114
|
+
|
|
1115
|
+
@over
|
|
1116
|
+
def __init__(self, value: int32) -> None:
|
|
1117
|
+
"""Construct a vector filled with a value."""
|
|
1118
|
+
...
|
|
1119
|
+
|
|
1120
|
+
class vec4ui:
|
|
1121
|
+
@over
|
|
1122
|
+
def __init__(self) -> None:
|
|
1123
|
+
"""Construct a zero-initialized vector."""
|
|
1124
|
+
...
|
|
1125
|
+
|
|
1126
|
+
@over
|
|
1127
|
+
def __init__(self, other: vec4ui) -> None:
|
|
1128
|
+
"""Construct a vector by copy."""
|
|
1129
|
+
...
|
|
1130
|
+
|
|
1131
|
+
@over
|
|
1132
|
+
def __init__(self, x: uint32, y: uint32, z: uint32, w: uint32) -> None:
|
|
1133
|
+
"""Construct a vector from its component values."""
|
|
1134
|
+
...
|
|
1135
|
+
|
|
1136
|
+
@over
|
|
1137
|
+
def __init__(self, args: Sequence[uint32]) -> None:
|
|
1138
|
+
"""Construct a vector from a sequence of values."""
|
|
1139
|
+
...
|
|
1140
|
+
|
|
1141
|
+
@over
|
|
1142
|
+
def __init__(self, value: uint32) -> None:
|
|
1143
|
+
"""Construct a vector filled with a value."""
|
|
1144
|
+
...
|
|
1145
|
+
|
|
1146
|
+
class vec4l:
|
|
1147
|
+
@over
|
|
1148
|
+
def __init__(self) -> None:
|
|
1149
|
+
"""Construct a zero-initialized vector."""
|
|
1150
|
+
...
|
|
1151
|
+
|
|
1152
|
+
@over
|
|
1153
|
+
def __init__(self, other: vec4l) -> None:
|
|
1154
|
+
"""Construct a vector by copy."""
|
|
1155
|
+
...
|
|
1156
|
+
|
|
1157
|
+
@over
|
|
1158
|
+
def __init__(self, x: int64, y: int64, z: int64, w: int64) -> None:
|
|
1159
|
+
"""Construct a vector from its component values."""
|
|
1160
|
+
...
|
|
1161
|
+
|
|
1162
|
+
@over
|
|
1163
|
+
def __init__(self, args: Sequence[int64]) -> None:
|
|
1164
|
+
"""Construct a vector from a sequence of values."""
|
|
1165
|
+
...
|
|
1166
|
+
|
|
1167
|
+
@over
|
|
1168
|
+
def __init__(self, value: int64) -> None:
|
|
1169
|
+
"""Construct a vector filled with a value."""
|
|
1170
|
+
...
|
|
1171
|
+
|
|
1172
|
+
class vec4ul:
|
|
1173
|
+
@over
|
|
1174
|
+
def __init__(self) -> None:
|
|
1175
|
+
"""Construct a zero-initialized vector."""
|
|
1176
|
+
...
|
|
1177
|
+
|
|
1178
|
+
@over
|
|
1179
|
+
def __init__(self, other: vec4ul) -> None:
|
|
1180
|
+
"""Construct a vector by copy."""
|
|
1181
|
+
...
|
|
1182
|
+
|
|
1183
|
+
@over
|
|
1184
|
+
def __init__(self, x: uint64, y: uint64, z: uint64, w: uint64) -> None:
|
|
1185
|
+
"""Construct a vector from its component values."""
|
|
1186
|
+
...
|
|
1187
|
+
|
|
1188
|
+
@over
|
|
1189
|
+
def __init__(self, args: Sequence[uint64]) -> None:
|
|
1190
|
+
"""Construct a vector from a sequence of values."""
|
|
1191
|
+
...
|
|
1192
|
+
|
|
1193
|
+
@over
|
|
1194
|
+
def __init__(self, value: uint64) -> None:
|
|
1195
|
+
"""Construct a vector filled with a value."""
|
|
1196
|
+
...
|
|
1197
|
+
|
|
1198
|
+
vec4 = vec4f
|
|
1199
|
+
|
|
1200
|
+
class mat22h:
|
|
1201
|
+
@over
|
|
1202
|
+
def __init__(self) -> None:
|
|
1203
|
+
"""Construct a zero-initialized matrix."""
|
|
1204
|
+
...
|
|
1205
|
+
|
|
1206
|
+
@over
|
|
1207
|
+
def __init__(self, other: mat22h) -> None:
|
|
1208
|
+
"""Construct a matrix by copy."""
|
|
1209
|
+
...
|
|
1210
|
+
|
|
1211
|
+
@over
|
|
1212
|
+
def __init__(self, m00: float16, m01: float16, m10: float16, m11: float16) -> None:
|
|
1213
|
+
"""Construct a matrix from its component values."""
|
|
1214
|
+
...
|
|
1215
|
+
|
|
1216
|
+
@over
|
|
1217
|
+
def __init__(self, v0: vec2h, v1: vec2h) -> None:
|
|
1218
|
+
"""Construct a matrix from its row vectors."""
|
|
1219
|
+
...
|
|
55
1220
|
|
|
56
|
-
|
|
1221
|
+
@over
|
|
1222
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
1223
|
+
"""Construct a matrix from a sequence of values."""
|
|
1224
|
+
...
|
|
57
1225
|
|
|
58
|
-
|
|
59
|
-
|
|
1226
|
+
@over
|
|
1227
|
+
def __init__(self, value: float16) -> None:
|
|
1228
|
+
"""Construct a matrix filled with a value."""
|
|
1229
|
+
...
|
|
60
1230
|
|
|
61
|
-
|
|
1231
|
+
class mat22f:
|
|
1232
|
+
@over
|
|
1233
|
+
def __init__(self) -> None:
|
|
1234
|
+
"""Construct a zero-initialized matrix."""
|
|
1235
|
+
...
|
|
62
1236
|
|
|
63
|
-
|
|
64
|
-
|
|
1237
|
+
@over
|
|
1238
|
+
def __init__(self, other: mat22f) -> None:
|
|
1239
|
+
"""Construct a matrix by copy."""
|
|
1240
|
+
...
|
|
65
1241
|
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
from
|
|
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
|
|
1242
|
+
@over
|
|
1243
|
+
def __init__(self, m00: float32, m01: float32, m10: float32, m11: float32) -> None:
|
|
1244
|
+
"""Construct a matrix from its component values."""
|
|
1245
|
+
...
|
|
144
1246
|
|
|
145
|
-
|
|
1247
|
+
@over
|
|
1248
|
+
def __init__(self, v0: vec2f, v1: vec2f) -> None:
|
|
1249
|
+
"""Construct a matrix from its row vectors."""
|
|
1250
|
+
...
|
|
146
1251
|
|
|
147
|
-
|
|
148
|
-
|
|
1252
|
+
@over
|
|
1253
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
1254
|
+
"""Construct a matrix from a sequence of values."""
|
|
1255
|
+
...
|
|
149
1256
|
|
|
150
|
-
|
|
1257
|
+
@over
|
|
1258
|
+
def __init__(self, value: float32) -> None:
|
|
1259
|
+
"""Construct a matrix filled with a value."""
|
|
1260
|
+
...
|
|
151
1261
|
|
|
152
|
-
|
|
1262
|
+
class mat22d:
|
|
1263
|
+
@over
|
|
1264
|
+
def __init__(self) -> None:
|
|
1265
|
+
"""Construct a zero-initialized matrix."""
|
|
1266
|
+
...
|
|
153
1267
|
|
|
154
|
-
|
|
1268
|
+
@over
|
|
1269
|
+
def __init__(self, other: mat22d) -> None:
|
|
1270
|
+
"""Construct a matrix by copy."""
|
|
1271
|
+
...
|
|
1272
|
+
|
|
1273
|
+
@over
|
|
1274
|
+
def __init__(self, m00: float64, m01: float64, m10: float64, m11: float64) -> None:
|
|
1275
|
+
"""Construct a matrix from its component values."""
|
|
1276
|
+
...
|
|
1277
|
+
|
|
1278
|
+
@over
|
|
1279
|
+
def __init__(self, v0: vec2d, v1: vec2d) -> None:
|
|
1280
|
+
"""Construct a matrix from its row vectors."""
|
|
1281
|
+
...
|
|
1282
|
+
|
|
1283
|
+
@over
|
|
1284
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
1285
|
+
"""Construct a matrix from a sequence of values."""
|
|
1286
|
+
...
|
|
1287
|
+
|
|
1288
|
+
@over
|
|
1289
|
+
def __init__(self, value: float64) -> None:
|
|
1290
|
+
"""Construct a matrix filled with a value."""
|
|
1291
|
+
...
|
|
1292
|
+
|
|
1293
|
+
mat22 = mat22f
|
|
1294
|
+
|
|
1295
|
+
class mat33h:
|
|
1296
|
+
@over
|
|
1297
|
+
def __init__(self) -> None:
|
|
1298
|
+
"""Construct a zero-initialized matrix."""
|
|
1299
|
+
...
|
|
1300
|
+
|
|
1301
|
+
@over
|
|
1302
|
+
def __init__(self, other: mat33h) -> None:
|
|
1303
|
+
"""Construct a matrix by copy."""
|
|
1304
|
+
...
|
|
1305
|
+
|
|
1306
|
+
@over
|
|
1307
|
+
def __init__(
|
|
1308
|
+
self,
|
|
1309
|
+
m00: float16,
|
|
1310
|
+
m01: float16,
|
|
1311
|
+
m02: float16,
|
|
1312
|
+
m10: float16,
|
|
1313
|
+
m11: float16,
|
|
1314
|
+
m12: float16,
|
|
1315
|
+
m20: float16,
|
|
1316
|
+
m21: float16,
|
|
1317
|
+
m22: float16,
|
|
1318
|
+
) -> None:
|
|
1319
|
+
"""Construct a matrix from its component values."""
|
|
1320
|
+
...
|
|
1321
|
+
|
|
1322
|
+
@over
|
|
1323
|
+
def __init__(self, v0: vec3h, v1: vec3h, v2: vec3h) -> None:
|
|
1324
|
+
"""Construct a matrix from its row vectors."""
|
|
1325
|
+
...
|
|
1326
|
+
|
|
1327
|
+
@over
|
|
1328
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
1329
|
+
"""Construct a matrix from a sequence of values."""
|
|
1330
|
+
...
|
|
1331
|
+
|
|
1332
|
+
@over
|
|
1333
|
+
def __init__(self, value: float16) -> None:
|
|
1334
|
+
"""Construct a matrix filled with a value."""
|
|
1335
|
+
...
|
|
1336
|
+
|
|
1337
|
+
class mat33f:
|
|
1338
|
+
@over
|
|
1339
|
+
def __init__(self) -> None:
|
|
1340
|
+
"""Construct a zero-initialized matrix."""
|
|
1341
|
+
...
|
|
1342
|
+
|
|
1343
|
+
@over
|
|
1344
|
+
def __init__(self, other: mat33f) -> None:
|
|
1345
|
+
"""Construct a matrix by copy."""
|
|
1346
|
+
...
|
|
1347
|
+
|
|
1348
|
+
@over
|
|
1349
|
+
def __init__(
|
|
1350
|
+
self,
|
|
1351
|
+
m00: float32,
|
|
1352
|
+
m01: float32,
|
|
1353
|
+
m02: float32,
|
|
1354
|
+
m10: float32,
|
|
1355
|
+
m11: float32,
|
|
1356
|
+
m12: float32,
|
|
1357
|
+
m20: float32,
|
|
1358
|
+
m21: float32,
|
|
1359
|
+
m22: float32,
|
|
1360
|
+
) -> None:
|
|
1361
|
+
"""Construct a matrix from its component values."""
|
|
1362
|
+
...
|
|
1363
|
+
|
|
1364
|
+
@over
|
|
1365
|
+
def __init__(self, v0: vec3f, v1: vec3f, v2: vec3f) -> None:
|
|
1366
|
+
"""Construct a matrix from its row vectors."""
|
|
1367
|
+
...
|
|
1368
|
+
|
|
1369
|
+
@over
|
|
1370
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
1371
|
+
"""Construct a matrix from a sequence of values."""
|
|
1372
|
+
...
|
|
1373
|
+
|
|
1374
|
+
@over
|
|
1375
|
+
def __init__(self, value: float32) -> None:
|
|
1376
|
+
"""Construct a matrix filled with a value."""
|
|
1377
|
+
...
|
|
1378
|
+
|
|
1379
|
+
class mat33d:
|
|
1380
|
+
@over
|
|
1381
|
+
def __init__(self) -> None:
|
|
1382
|
+
"""Construct a zero-initialized matrix."""
|
|
1383
|
+
...
|
|
1384
|
+
|
|
1385
|
+
@over
|
|
1386
|
+
def __init__(self, other: mat33d) -> None:
|
|
1387
|
+
"""Construct a matrix by copy."""
|
|
1388
|
+
...
|
|
1389
|
+
|
|
1390
|
+
@over
|
|
1391
|
+
def __init__(
|
|
1392
|
+
self,
|
|
1393
|
+
m00: float64,
|
|
1394
|
+
m01: float64,
|
|
1395
|
+
m02: float64,
|
|
1396
|
+
m10: float64,
|
|
1397
|
+
m11: float64,
|
|
1398
|
+
m12: float64,
|
|
1399
|
+
m20: float64,
|
|
1400
|
+
m21: float64,
|
|
1401
|
+
m22: float64,
|
|
1402
|
+
) -> None:
|
|
1403
|
+
"""Construct a matrix from its component values."""
|
|
1404
|
+
...
|
|
1405
|
+
|
|
1406
|
+
@over
|
|
1407
|
+
def __init__(self, v0: vec3d, v1: vec3d, v2: vec3d) -> None:
|
|
1408
|
+
"""Construct a matrix from its row vectors."""
|
|
1409
|
+
...
|
|
1410
|
+
|
|
1411
|
+
@over
|
|
1412
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
1413
|
+
"""Construct a matrix from a sequence of values."""
|
|
1414
|
+
...
|
|
1415
|
+
|
|
1416
|
+
@over
|
|
1417
|
+
def __init__(self, value: float64) -> None:
|
|
1418
|
+
"""Construct a matrix filled with a value."""
|
|
1419
|
+
...
|
|
1420
|
+
|
|
1421
|
+
mat33 = mat33f
|
|
1422
|
+
|
|
1423
|
+
class mat44h:
|
|
1424
|
+
@over
|
|
1425
|
+
def __init__(self) -> None:
|
|
1426
|
+
"""Construct a zero-initialized matrix."""
|
|
1427
|
+
...
|
|
1428
|
+
|
|
1429
|
+
@over
|
|
1430
|
+
def __init__(self, other: mat44h) -> None:
|
|
1431
|
+
"""Construct a matrix by copy."""
|
|
1432
|
+
...
|
|
1433
|
+
|
|
1434
|
+
@over
|
|
1435
|
+
def __init__(
|
|
1436
|
+
self,
|
|
1437
|
+
m00: float16,
|
|
1438
|
+
m01: float16,
|
|
1439
|
+
m02: float16,
|
|
1440
|
+
m03: float16,
|
|
1441
|
+
m10: float16,
|
|
1442
|
+
m11: float16,
|
|
1443
|
+
m12: float16,
|
|
1444
|
+
m13: float16,
|
|
1445
|
+
m20: float16,
|
|
1446
|
+
m21: float16,
|
|
1447
|
+
m22: float16,
|
|
1448
|
+
m23: float16,
|
|
1449
|
+
m30: float16,
|
|
1450
|
+
m31: float16,
|
|
1451
|
+
m32: float16,
|
|
1452
|
+
m33: float16,
|
|
1453
|
+
) -> None:
|
|
1454
|
+
"""Construct a matrix from its component values."""
|
|
1455
|
+
...
|
|
1456
|
+
|
|
1457
|
+
@over
|
|
1458
|
+
def __init__(self, v0: vec4h, v1: vec4h, v2: vec4h, v3: vec4h) -> None:
|
|
1459
|
+
"""Construct a matrix from its row vectors."""
|
|
1460
|
+
...
|
|
1461
|
+
|
|
1462
|
+
@over
|
|
1463
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
1464
|
+
"""Construct a matrix from a sequence of values."""
|
|
1465
|
+
...
|
|
1466
|
+
|
|
1467
|
+
@over
|
|
1468
|
+
def __init__(self, value: float16) -> None:
|
|
1469
|
+
"""Construct a matrix filled with a value."""
|
|
1470
|
+
...
|
|
1471
|
+
|
|
1472
|
+
class mat44f:
|
|
1473
|
+
@over
|
|
1474
|
+
def __init__(self) -> None:
|
|
1475
|
+
"""Construct a zero-initialized matrix."""
|
|
1476
|
+
...
|
|
1477
|
+
|
|
1478
|
+
@over
|
|
1479
|
+
def __init__(self, other: mat44f) -> None:
|
|
1480
|
+
"""Construct a matrix by copy."""
|
|
1481
|
+
...
|
|
1482
|
+
|
|
1483
|
+
@over
|
|
1484
|
+
def __init__(
|
|
1485
|
+
self,
|
|
1486
|
+
m00: float32,
|
|
1487
|
+
m01: float32,
|
|
1488
|
+
m02: float32,
|
|
1489
|
+
m03: float32,
|
|
1490
|
+
m10: float32,
|
|
1491
|
+
m11: float32,
|
|
1492
|
+
m12: float32,
|
|
1493
|
+
m13: float32,
|
|
1494
|
+
m20: float32,
|
|
1495
|
+
m21: float32,
|
|
1496
|
+
m22: float32,
|
|
1497
|
+
m23: float32,
|
|
1498
|
+
m30: float32,
|
|
1499
|
+
m31: float32,
|
|
1500
|
+
m32: float32,
|
|
1501
|
+
m33: float32,
|
|
1502
|
+
) -> None:
|
|
1503
|
+
"""Construct a matrix from its component values."""
|
|
1504
|
+
...
|
|
1505
|
+
|
|
1506
|
+
@over
|
|
1507
|
+
def __init__(self, v0: vec4f, v1: vec4f, v2: vec4f, v3: vec4f) -> None:
|
|
1508
|
+
"""Construct a matrix from its row vectors."""
|
|
1509
|
+
...
|
|
1510
|
+
|
|
1511
|
+
@over
|
|
1512
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
1513
|
+
"""Construct a matrix from a sequence of values."""
|
|
1514
|
+
...
|
|
1515
|
+
|
|
1516
|
+
@over
|
|
1517
|
+
def __init__(self, value: float32) -> None:
|
|
1518
|
+
"""Construct a matrix filled with a value."""
|
|
1519
|
+
...
|
|
1520
|
+
|
|
1521
|
+
class mat44d:
|
|
1522
|
+
@over
|
|
1523
|
+
def __init__(self) -> None:
|
|
1524
|
+
"""Construct a zero-initialized matrix."""
|
|
1525
|
+
...
|
|
1526
|
+
|
|
1527
|
+
@over
|
|
1528
|
+
def __init__(self, other: mat44d) -> None:
|
|
1529
|
+
"""Construct a matrix by copy."""
|
|
1530
|
+
...
|
|
1531
|
+
|
|
1532
|
+
@over
|
|
1533
|
+
def __init__(
|
|
1534
|
+
self,
|
|
1535
|
+
m00: float64,
|
|
1536
|
+
m01: float64,
|
|
1537
|
+
m02: float64,
|
|
1538
|
+
m03: float64,
|
|
1539
|
+
m10: float64,
|
|
1540
|
+
m11: float64,
|
|
1541
|
+
m12: float64,
|
|
1542
|
+
m13: float64,
|
|
1543
|
+
m20: float64,
|
|
1544
|
+
m21: float64,
|
|
1545
|
+
m22: float64,
|
|
1546
|
+
m23: float64,
|
|
1547
|
+
m30: float64,
|
|
1548
|
+
m31: float64,
|
|
1549
|
+
m32: float64,
|
|
1550
|
+
m33: float64,
|
|
1551
|
+
) -> None:
|
|
1552
|
+
"""Construct a matrix from its component values."""
|
|
1553
|
+
...
|
|
1554
|
+
|
|
1555
|
+
@over
|
|
1556
|
+
def __init__(self, v0: vec4d, v1: vec4d, v2: vec4d, v3: vec4d) -> None:
|
|
1557
|
+
"""Construct a matrix from its row vectors."""
|
|
1558
|
+
...
|
|
1559
|
+
|
|
1560
|
+
@over
|
|
1561
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
1562
|
+
"""Construct a matrix from a sequence of values."""
|
|
1563
|
+
...
|
|
1564
|
+
|
|
1565
|
+
@over
|
|
1566
|
+
def __init__(self, value: float64) -> None:
|
|
1567
|
+
"""Construct a matrix filled with a value."""
|
|
1568
|
+
...
|
|
1569
|
+
|
|
1570
|
+
mat44 = mat44f
|
|
1571
|
+
|
|
1572
|
+
class quath:
|
|
1573
|
+
@over
|
|
1574
|
+
def __init__(self) -> None:
|
|
1575
|
+
"""Construct a zero-initialized quaternion."""
|
|
1576
|
+
...
|
|
1577
|
+
|
|
1578
|
+
@over
|
|
1579
|
+
def __init__(self, other: quath) -> None:
|
|
1580
|
+
"""Construct a quaternion by copy."""
|
|
1581
|
+
...
|
|
1582
|
+
|
|
1583
|
+
@over
|
|
1584
|
+
def __init__(self, x: float16, y: float16, z: float16, w: float16) -> None:
|
|
1585
|
+
"""Construct a quaternion from its component values."""
|
|
1586
|
+
...
|
|
1587
|
+
|
|
1588
|
+
@over
|
|
1589
|
+
def __init__(self, args: Sequence[float16]) -> None:
|
|
1590
|
+
"""Construct a quaternion from a sequence of values."""
|
|
1591
|
+
...
|
|
1592
|
+
|
|
1593
|
+
@over
|
|
1594
|
+
def __init__(self, value: float16) -> None:
|
|
1595
|
+
"""Construct a quaternion filled with a value."""
|
|
1596
|
+
...
|
|
1597
|
+
|
|
1598
|
+
class quatf:
|
|
1599
|
+
@over
|
|
1600
|
+
def __init__(self) -> None:
|
|
1601
|
+
"""Construct a zero-initialized quaternion."""
|
|
1602
|
+
...
|
|
1603
|
+
|
|
1604
|
+
@over
|
|
1605
|
+
def __init__(self, other: quatf) -> None:
|
|
1606
|
+
"""Construct a quaternion by copy."""
|
|
1607
|
+
...
|
|
1608
|
+
|
|
1609
|
+
@over
|
|
1610
|
+
def __init__(self, x: float32, y: float32, z: float32, w: float32) -> None:
|
|
1611
|
+
"""Construct a quaternion from its component values."""
|
|
1612
|
+
...
|
|
1613
|
+
|
|
1614
|
+
@over
|
|
1615
|
+
def __init__(self, args: Sequence[float32]) -> None:
|
|
1616
|
+
"""Construct a quaternion from a sequence of values."""
|
|
1617
|
+
...
|
|
1618
|
+
|
|
1619
|
+
@over
|
|
1620
|
+
def __init__(self, value: float32) -> None:
|
|
1621
|
+
"""Construct a quaternion filled with a value."""
|
|
1622
|
+
...
|
|
1623
|
+
|
|
1624
|
+
class quatd:
|
|
1625
|
+
@over
|
|
1626
|
+
def __init__(self) -> None:
|
|
1627
|
+
"""Construct a zero-initialized quaternion."""
|
|
1628
|
+
...
|
|
1629
|
+
|
|
1630
|
+
@over
|
|
1631
|
+
def __init__(self, other: quatd) -> None:
|
|
1632
|
+
"""Construct a quaternion by copy."""
|
|
1633
|
+
...
|
|
1634
|
+
|
|
1635
|
+
@over
|
|
1636
|
+
def __init__(self, x: float64, y: float64, z: float64, w: float64) -> None:
|
|
1637
|
+
"""Construct a quaternion from its component values."""
|
|
1638
|
+
...
|
|
1639
|
+
|
|
1640
|
+
@over
|
|
1641
|
+
def __init__(self, args: Sequence[float64]) -> None:
|
|
1642
|
+
"""Construct a quaternion from a sequence of values."""
|
|
1643
|
+
...
|
|
1644
|
+
|
|
1645
|
+
@over
|
|
1646
|
+
def __init__(self, value: float64) -> None:
|
|
1647
|
+
"""Construct a quaternion filled with a value."""
|
|
1648
|
+
...
|
|
1649
|
+
|
|
1650
|
+
quat = quatf
|
|
1651
|
+
|
|
1652
|
+
class transformh:
|
|
1653
|
+
@over
|
|
1654
|
+
def __init__(self) -> None:
|
|
1655
|
+
"""Construct a zero-initialized transformation."""
|
|
1656
|
+
...
|
|
1657
|
+
|
|
1658
|
+
@over
|
|
1659
|
+
def __init__(self, other: transformh) -> None:
|
|
1660
|
+
"""Construct a transformation by copy."""
|
|
1661
|
+
...
|
|
1662
|
+
|
|
1663
|
+
@over
|
|
1664
|
+
def __init__(self, p: vec3h, q: quath) -> None:
|
|
1665
|
+
"""Construct a transformation from its p and q components."""
|
|
1666
|
+
...
|
|
1667
|
+
|
|
1668
|
+
@over
|
|
1669
|
+
def __init__(
|
|
1670
|
+
self, px: float16, py: float16, pz: float16, qx: float16, qy: float16, qz: float16, qw: float16
|
|
1671
|
+
) -> None:
|
|
1672
|
+
"""Construct a transformation from its component values."""
|
|
1673
|
+
...
|
|
1674
|
+
|
|
1675
|
+
@over
|
|
1676
|
+
def __init__(self, p: Sequence[float16], q: Sequence[float16]) -> None:
|
|
1677
|
+
"""Construct a transformation from two sequences of values."""
|
|
1678
|
+
...
|
|
1679
|
+
|
|
1680
|
+
@over
|
|
1681
|
+
def __init__(self, value: float16) -> None:
|
|
1682
|
+
"""Construct a transformation filled with a value."""
|
|
1683
|
+
...
|
|
1684
|
+
|
|
1685
|
+
class transformf:
|
|
1686
|
+
@over
|
|
1687
|
+
def __init__(self) -> None:
|
|
1688
|
+
"""Construct a zero-initialized transformation."""
|
|
1689
|
+
...
|
|
1690
|
+
|
|
1691
|
+
@over
|
|
1692
|
+
def __init__(self, other: transformf) -> None:
|
|
1693
|
+
"""Construct a transformation by copy."""
|
|
1694
|
+
...
|
|
1695
|
+
|
|
1696
|
+
@over
|
|
1697
|
+
def __init__(self, p: vec3f, q: quatf) -> None:
|
|
1698
|
+
"""Construct a transformation from its p and q components."""
|
|
1699
|
+
...
|
|
1700
|
+
|
|
1701
|
+
@over
|
|
1702
|
+
def __init__(
|
|
1703
|
+
self, px: float32, py: float32, pz: float32, qx: float32, qy: float32, qz: float32, qw: float32
|
|
1704
|
+
) -> None:
|
|
1705
|
+
"""Construct a transformation from its component values."""
|
|
1706
|
+
...
|
|
1707
|
+
|
|
1708
|
+
@over
|
|
1709
|
+
def __init__(self, p: Sequence[float32], q: Sequence[float32]) -> None:
|
|
1710
|
+
"""Construct a transformation from two sequences of values."""
|
|
1711
|
+
...
|
|
1712
|
+
|
|
1713
|
+
@over
|
|
1714
|
+
def __init__(self, value: float32) -> None:
|
|
1715
|
+
"""Construct a transformation filled with a value."""
|
|
1716
|
+
...
|
|
1717
|
+
|
|
1718
|
+
class transformd:
|
|
1719
|
+
@over
|
|
1720
|
+
def __init__(self) -> None:
|
|
1721
|
+
"""Construct a zero-initialized transformation."""
|
|
1722
|
+
...
|
|
1723
|
+
|
|
1724
|
+
@over
|
|
1725
|
+
def __init__(self, other: transformd) -> None:
|
|
1726
|
+
"""Construct a transformation by copy."""
|
|
1727
|
+
...
|
|
1728
|
+
|
|
1729
|
+
@over
|
|
1730
|
+
def __init__(self, p: vec3d, q: quatd) -> None:
|
|
1731
|
+
"""Construct a transformation from its p and q components."""
|
|
1732
|
+
...
|
|
1733
|
+
|
|
1734
|
+
@over
|
|
1735
|
+
def __init__(
|
|
1736
|
+
self, px: float64, py: float64, pz: float64, qx: float64, qy: float64, qz: float64, qw: float64
|
|
1737
|
+
) -> None:
|
|
1738
|
+
"""Construct a transformation from its component values."""
|
|
1739
|
+
...
|
|
1740
|
+
|
|
1741
|
+
@over
|
|
1742
|
+
def __init__(self, p: Sequence[float64], q: Sequence[float64]) -> None:
|
|
1743
|
+
"""Construct a transformation from two sequences of values."""
|
|
1744
|
+
...
|
|
1745
|
+
|
|
1746
|
+
@over
|
|
1747
|
+
def __init__(self, value: float64) -> None:
|
|
1748
|
+
"""Construct a transformation filled with a value."""
|
|
1749
|
+
...
|
|
1750
|
+
|
|
1751
|
+
transform = transformf
|
|
155
1752
|
|
|
156
1753
|
@over
|
|
157
1754
|
def min(a: Scalar, b: Scalar) -> Scalar:
|
|
@@ -924,7 +2521,7 @@ def tile_arange(*args: Scalar, dtype: Scalar, storage: str) -> Tile[Scalar, Tupl
|
|
|
924
2521
|
|
|
925
2522
|
@over
|
|
926
2523
|
def tile_load(
|
|
927
|
-
a: Array[Any], shape: Tuple[int, ...], offset: Tuple[int, ...], storage: str
|
|
2524
|
+
a: Array[Any], shape: Tuple[int, ...], offset: Tuple[int, ...], storage: str, bounds_check: bool
|
|
928
2525
|
) -> Tile[Any, Tuple[int, ...]]:
|
|
929
2526
|
"""Loads a tile from a global memory array.
|
|
930
2527
|
|
|
@@ -935,12 +2532,80 @@ def tile_load(
|
|
|
935
2532
|
:param offset: Offset in the source array to begin reading from (optional)
|
|
936
2533
|
:param storage: The storage location for the tile: ``"register"`` for registers
|
|
937
2534
|
(default) or ``"shared"`` for shared memory.
|
|
2535
|
+
:param bounds_check: Needed for unaligned tiles, but can disable for memory-aligned tiles for faster load times
|
|
2536
|
+
:returns: A tile with shape as specified and data type the same as the source array
|
|
2537
|
+
"""
|
|
2538
|
+
...
|
|
2539
|
+
|
|
2540
|
+
@over
|
|
2541
|
+
def tile_load_indexed(
|
|
2542
|
+
a: Array[Any],
|
|
2543
|
+
indices: Tile[int32, Tuple[int]],
|
|
2544
|
+
shape: Tuple[int, ...],
|
|
2545
|
+
offset: Tuple[int, ...],
|
|
2546
|
+
axis: int32,
|
|
2547
|
+
storage: str,
|
|
2548
|
+
) -> Tile[Any, Tuple[int, ...]]:
|
|
2549
|
+
"""Loads a tile from a global memory array, with loads along a specified axis mapped according to a 1D tile of indices.
|
|
2550
|
+
|
|
2551
|
+
:param a: The source array in global memory
|
|
2552
|
+
:param indices: A 1D tile of integer indices mapping to elements in ``a``.
|
|
2553
|
+
:param shape: Shape of the tile to load, must have the same number of dimensions as ``a``, and along ``axis``, it must have the same number of elements as the ``indices`` tile.
|
|
2554
|
+
:param offset: Offset in the source array to begin reading from (optional)
|
|
2555
|
+
:param axis: Axis of ``a`` that indices refer to
|
|
2556
|
+
:param storage: The storage location for the tile: ``"register"`` for registers (default) or ``"shared"`` for shared memory.
|
|
938
2557
|
:returns: A tile with shape as specified and data type the same as the source array
|
|
2558
|
+
|
|
2559
|
+
This example shows how to select and store the even indexed rows from a 2D array:
|
|
2560
|
+
|
|
2561
|
+
.. code-block:: python
|
|
2562
|
+
|
|
2563
|
+
TILE_M = wp.constant(2)
|
|
2564
|
+
TILE_N = wp.constant(2)
|
|
2565
|
+
HALF_M = wp.constant(TILE_M // 2)
|
|
2566
|
+
HALF_N = wp.constant(TILE_N // 2)
|
|
2567
|
+
|
|
2568
|
+
@wp.kernel
|
|
2569
|
+
def compute(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
|
|
2570
|
+
i, j = wp.tid()
|
|
2571
|
+
|
|
2572
|
+
evens = wp.tile_arange(HALF_M, dtype=int, storage="shared") * 2
|
|
2573
|
+
|
|
2574
|
+
t0 = wp.tile_load_indexed(
|
|
2575
|
+
x, indices=evens, shape=(HALF_M, TILE_N), offset=(i * TILE_M, j * TILE_N), axis=0, storage="register"
|
|
2576
|
+
)
|
|
2577
|
+
wp.tile_store(y, t0, offset=(i * HALF_M, j * TILE_N))
|
|
2578
|
+
|
|
2579
|
+
M = TILE_M * 2
|
|
2580
|
+
N = TILE_N * 2
|
|
2581
|
+
|
|
2582
|
+
arr = np.arange(M * N).reshape(M, N)
|
|
2583
|
+
|
|
2584
|
+
x = wp.array(arr, dtype=float)
|
|
2585
|
+
y = wp.zeros((M // 2, N), dtype=float)
|
|
2586
|
+
|
|
2587
|
+
wp.launch_tiled(compute, dim=[2, 2], inputs=[x], outputs=[y], block_dim=32, device=device)
|
|
2588
|
+
|
|
2589
|
+
print(x.numpy())
|
|
2590
|
+
print(y.numpy())
|
|
2591
|
+
|
|
2592
|
+
Prints:
|
|
2593
|
+
|
|
2594
|
+
.. code-block:: text
|
|
2595
|
+
|
|
2596
|
+
[[ 0. 1. 2. 3.]
|
|
2597
|
+
[ 4. 5. 6. 7.]
|
|
2598
|
+
[ 8. 9. 10. 11.]
|
|
2599
|
+
[12. 13. 14. 15.]]
|
|
2600
|
+
|
|
2601
|
+
[[ 0. 1. 2. 3.]
|
|
2602
|
+
[ 8. 9. 10. 11.]]
|
|
2603
|
+
|
|
939
2604
|
"""
|
|
940
2605
|
...
|
|
941
2606
|
|
|
942
2607
|
@over
|
|
943
|
-
def tile_store(a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...]):
|
|
2608
|
+
def tile_store(a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...], bounds_check: bool):
|
|
944
2609
|
"""Store a tile to a global memory array.
|
|
945
2610
|
|
|
946
2611
|
This method will cooperatively store a tile to global memory using all threads in the block.
|
|
@@ -948,22 +2613,147 @@ def tile_store(a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int,
|
|
|
948
2613
|
:param a: The destination array in global memory
|
|
949
2614
|
:param t: The source tile to store data from, must have the same data type and number of dimensions as the destination array
|
|
950
2615
|
:param offset: Offset in the destination array (optional)
|
|
2616
|
+
:param bounds_check: Needed for unaligned tiles, but can disable for memory-aligned tiles for faster write times
|
|
2617
|
+
|
|
2618
|
+
"""
|
|
2619
|
+
...
|
|
2620
|
+
|
|
2621
|
+
@over
|
|
2622
|
+
def tile_store_indexed(
|
|
2623
|
+
a: Array[Any], indices: Tile[int32, Tuple[int]], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...], axis: int32
|
|
2624
|
+
):
|
|
2625
|
+
"""Store a tile to a global memory array, with storage along a specified axis mapped according to a 1D tile of indices.
|
|
2626
|
+
|
|
2627
|
+
:param a: The destination array in global memory
|
|
2628
|
+
:param indices: A 1D tile of integer indices mapping to elements in ``a``.
|
|
2629
|
+
:param t: The source tile to store data from, must have the same data type and number of dimensions as the destination array, and along ``axis``, it must have the same number of elements as the ``indices`` tile.
|
|
2630
|
+
:param offset: Offset in the destination array (optional)
|
|
2631
|
+
:param axis: Axis of ``a`` that indices refer to
|
|
2632
|
+
|
|
2633
|
+
This example shows how to map tile rows to the even rows of a 2D array:
|
|
2634
|
+
|
|
2635
|
+
.. code-block:: python
|
|
2636
|
+
|
|
2637
|
+
TILE_M = wp.constant(2)
|
|
2638
|
+
TILE_N = wp.constant(2)
|
|
2639
|
+
TWO_M = wp.constant(TILE_M * 2)
|
|
2640
|
+
TWO_N = wp.constant(TILE_N * 2)
|
|
2641
|
+
|
|
2642
|
+
@wp.kernel
|
|
2643
|
+
def compute(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
|
|
2644
|
+
i, j = wp.tid()
|
|
2645
|
+
|
|
2646
|
+
t = wp.tile_load(x, shape=(TILE_M, TILE_N), offset=(i * TILE_M, j * TILE_N), storage="register")
|
|
2647
|
+
|
|
2648
|
+
evens_M = wp.tile_arange(TILE_M, dtype=int, storage="shared") * 2
|
|
2649
|
+
|
|
2650
|
+
wp.tile_store_indexed(y, indices=evens_M, t=t, offset=(i * TWO_M, j * TILE_N), axis=0)
|
|
2651
|
+
|
|
2652
|
+
M = TILE_M * 2
|
|
2653
|
+
N = TILE_N * 2
|
|
2654
|
+
|
|
2655
|
+
arr = np.arange(M * N, dtype=float).reshape(M, N)
|
|
2656
|
+
|
|
2657
|
+
x = wp.array(arr, dtype=float, requires_grad=True, device=device)
|
|
2658
|
+
y = wp.zeros((M * 2, N), dtype=float, requires_grad=True, device=device)
|
|
2659
|
+
|
|
2660
|
+
wp.launch_tiled(compute, dim=[2, 2], inputs=[x], outputs=[y], block_dim=32, device=device)
|
|
2661
|
+
|
|
2662
|
+
print(x.numpy())
|
|
2663
|
+
print(y.numpy())
|
|
2664
|
+
|
|
2665
|
+
Prints:
|
|
2666
|
+
|
|
2667
|
+
.. code-block:: text
|
|
2668
|
+
|
|
2669
|
+
[[ 0. 1. 2. 3.]
|
|
2670
|
+
[ 4. 5. 6. 7.]
|
|
2671
|
+
[ 8. 9. 10. 11.]
|
|
2672
|
+
[12. 13. 14. 15.]]
|
|
2673
|
+
|
|
2674
|
+
[[ 0. 1. 2. 3.]
|
|
2675
|
+
[ 0. 0. 0. 0.]
|
|
2676
|
+
[ 4. 5. 6. 7.]
|
|
2677
|
+
[ 0. 0. 0. 0.]
|
|
2678
|
+
[ 8. 9. 10. 11.]
|
|
2679
|
+
[ 0. 0. 0. 0.]
|
|
2680
|
+
[12. 13. 14. 15.]
|
|
2681
|
+
[ 0. 0. 0. 0.]]
|
|
2682
|
+
|
|
951
2683
|
"""
|
|
952
2684
|
...
|
|
953
2685
|
|
|
954
2686
|
@over
|
|
955
2687
|
def tile_atomic_add(
|
|
956
|
-
a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...]
|
|
2688
|
+
a: Array[Any], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...], bounds_check: bool
|
|
957
2689
|
) -> Tile[Any, Tuple[int, ...]]:
|
|
958
2690
|
"""Atomically add a tile onto the array `a`, each element will be updated atomically.
|
|
959
2691
|
|
|
960
2692
|
:param a: Array in global memory, should have the same ``dtype`` as the input tile
|
|
961
2693
|
:param t: Source tile to add to the destination array
|
|
962
2694
|
:param offset: Offset in the destination array (optional)
|
|
2695
|
+
:param bounds_check: Needed for unaligned tiles, but can disable for memory-aligned tiles for faster write times
|
|
963
2696
|
:returns: A tile with the same dimensions and data type as the source tile, holding the original value of the destination elements
|
|
964
2697
|
"""
|
|
965
2698
|
...
|
|
966
2699
|
|
|
2700
|
+
@over
|
|
2701
|
+
def tile_atomic_add_indexed(
|
|
2702
|
+
a: Array[Any], indices: Tile[int32, Tuple[int]], t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...], axis: int32
|
|
2703
|
+
) -> Tile[Any, Tuple[int, ...]]:
|
|
2704
|
+
"""Atomically add a tile to a global memory array, with storage along a specified axis mapped according to a 1D tile of indices.
|
|
2705
|
+
|
|
2706
|
+
:param a: The destination array in global memory
|
|
2707
|
+
:param indices: A 1D tile of integer indices mapping to elements in ``a``.
|
|
2708
|
+
:param t: The source tile to extract data from, must have the same data type and number of dimensions as the destination array, and along ``axis``, it must have the same number of elements as the ``indices`` tile.
|
|
2709
|
+
:param offset: Offset in the destination array (optional)
|
|
2710
|
+
:param axis: Axis of ``a`` that indices refer to
|
|
2711
|
+
|
|
2712
|
+
This example shows how to compute a blocked, row-wise reduction:
|
|
2713
|
+
|
|
2714
|
+
.. code-block:: python
|
|
2715
|
+
|
|
2716
|
+
TILE_M = wp.constant(2)
|
|
2717
|
+
TILE_N = wp.constant(2)
|
|
2718
|
+
|
|
2719
|
+
@wp.kernel
|
|
2720
|
+
def tile_atomic_add_indexed(x: wp.array2d(dtype=float), y: wp.array2d(dtype=float)):
|
|
2721
|
+
i, j = wp.tid()
|
|
2722
|
+
|
|
2723
|
+
t = wp.tile_load(x, shape=(TILE_M, TILE_N), offset=(i * TILE_M, j * TILE_N), storage="register")
|
|
2724
|
+
|
|
2725
|
+
zeros = wp.tile_zeros(TILE_M, dtype=int, storage="shared")
|
|
2726
|
+
|
|
2727
|
+
wp.tile_atomic_add_indexed(y, indices=zeros, t=t, offset=(i, j * TILE_N), axis=0)
|
|
2728
|
+
|
|
2729
|
+
M = TILE_M * 2
|
|
2730
|
+
N = TILE_N * 2
|
|
2731
|
+
|
|
2732
|
+
arr = np.arange(M * N, dtype=float).reshape(M, N)
|
|
2733
|
+
|
|
2734
|
+
x = wp.array(arr, dtype=float, requires_grad=True, device=device)
|
|
2735
|
+
y = wp.zeros((2, N), dtype=float, requires_grad=True, device=device)
|
|
2736
|
+
|
|
2737
|
+
wp.launch_tiled(tile_atomic_add_indexed, dim=[2, 2], inputs=[x], outputs=[y], block_dim=32, device=device)
|
|
2738
|
+
|
|
2739
|
+
print(x.numpy())
|
|
2740
|
+
print(y.numpy())
|
|
2741
|
+
|
|
2742
|
+
Prints:
|
|
2743
|
+
|
|
2744
|
+
.. code-block:: text
|
|
2745
|
+
|
|
2746
|
+
[[ 0. 1. 2. 3.]
|
|
2747
|
+
[ 4. 5. 6. 7.]
|
|
2748
|
+
[ 8. 9. 10. 11.]
|
|
2749
|
+
[12. 13. 14. 15.]]
|
|
2750
|
+
|
|
2751
|
+
[[ 4. 6. 8. 10.]
|
|
2752
|
+
[20. 22. 24. 26.]]
|
|
2753
|
+
|
|
2754
|
+
"""
|
|
2755
|
+
...
|
|
2756
|
+
|
|
967
2757
|
@over
|
|
968
2758
|
def tile_view(
|
|
969
2759
|
t: Tile[Any, Tuple[int, ...]], offset: Tuple[int, ...], shape: Tuple[int, ...]
|
|
@@ -1370,7 +3160,7 @@ def tile_map(op: Callable, a: Tile[Scalar, Tuple[int, ...]]) -> Tile[Scalar, Tup
|
|
|
1370
3160
|
|
|
1371
3161
|
:param op: A callable function that accepts one argument and returns one argument, may be a user function or builtin
|
|
1372
3162
|
: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
|
|
3163
|
+
:returns: A tile with the same dimensions as the input tile. Its datatype is specified by the return type of op
|
|
1374
3164
|
|
|
1375
3165
|
Example:
|
|
1376
3166
|
|
|
@@ -1401,12 +3191,12 @@ def tile_map(
|
|
|
1401
3191
|
"""Apply a binary function onto the tile.
|
|
1402
3192
|
|
|
1403
3193
|
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
|
|
3194
|
+
Both input tiles must have the same dimensions, and if using a builtin op, the same datatypes.
|
|
1405
3195
|
|
|
1406
3196
|
: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
3197
|
:param a: The first input tile, the operator (or one of its overloads) must be able to accept the tile's dtype
|
|
1408
3198
|
: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
|
|
3199
|
+
:returns: A tile with the same dimensions as the input tiles. Its datatype is specified by the return type of op
|
|
1410
3200
|
|
|
1411
3201
|
Example:
|
|
1412
3202
|
|
|
@@ -1551,7 +3341,7 @@ def mesh_query_ray(id: uint64, start: vec3f, dir: vec3f, max_t: float32) -> Mesh
|
|
|
1551
3341
|
def mesh_query_aabb(id: uint64, low: vec3f, high: vec3f) -> MeshQueryAABB:
|
|
1552
3342
|
"""Construct an axis-aligned bounding box query against a :class:`Mesh`.
|
|
1553
3343
|
|
|
1554
|
-
This query can be used to iterate over all triangles inside a volume.
|
|
3344
|
+
This query can be used to iterate over all bounding boxes of the triangles inside a volume.
|
|
1555
3345
|
|
|
1556
3346
|
:param id: The mesh identifier
|
|
1557
3347
|
:param low: The lower bound of the bounding box in mesh space
|
|
@@ -1561,7 +3351,7 @@ def mesh_query_aabb(id: uint64, low: vec3f, high: vec3f) -> MeshQueryAABB:
|
|
|
1561
3351
|
|
|
1562
3352
|
@over
|
|
1563
3353
|
def mesh_query_aabb_next(query: MeshQueryAABB, index: int32) -> bool:
|
|
1564
|
-
"""Move to the next triangle
|
|
3354
|
+
"""Move to the next triangle whose bounding box overlaps the query bounding box.
|
|
1565
3355
|
|
|
1566
3356
|
The index of the current face is stored in ``index``, returns ``False`` if there are no more overlapping triangles.
|
|
1567
3357
|
"""
|
|
@@ -2971,7 +4761,7 @@ def mod(a: Scalar, b: Scalar) -> Scalar:
|
|
|
2971
4761
|
...
|
|
2972
4762
|
|
|
2973
4763
|
@over
|
|
2974
|
-
def mod(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Scalar:
|
|
4764
|
+
def mod(a: Vector[Any, Scalar], b: Vector[Any, Scalar]) -> Vector[Any, Scalar]:
|
|
2975
4765
|
"""Modulo operation using truncated division."""
|
|
2976
4766
|
...
|
|
2977
4767
|
|