warp-lang 1.8.0__py3-none-manylinux_2_34_aarch64.whl → 1.9.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 +282 -103
- warp/__init__.pyi +482 -110
- warp/bin/warp-clang.so +0 -0
- warp/bin/warp.so +0 -0
- warp/build.py +93 -30
- warp/build_dll.py +48 -63
- warp/builtins.py +955 -137
- warp/codegen.py +327 -209
- warp/config.py +1 -1
- warp/context.py +1363 -800
- 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_callable.py +34 -4
- warp/examples/interop/example_jax_kernel.py +27 -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 +266 -166
- warp/fem/geometry/geometry.py +5 -5
- warp/fem/integrate.py +200 -91
- warp/fem/space/restriction.py +4 -0
- warp/fem/space/shape/tet_shape_function.py +3 -10
- warp/jax_experimental/custom_call.py +1 -1
- warp/jax_experimental/ffi.py +203 -54
- warp/marching_cubes.py +708 -0
- warp/native/array.h +103 -8
- warp/native/builtin.h +90 -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 +13 -3
- warp/native/crt.cpp +2 -2
- warp/native/crt.h +3 -5
- warp/native/cuda_util.cpp +42 -11
- 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 +4 -4
- warp/native/mat.h +1913 -119
- warp/native/mathdx.cpp +43 -43
- warp/native/mesh.cpp +24 -24
- warp/native/mesh.cu +26 -26
- warp/native/mesh.h +5 -3
- 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 +337 -16
- warp/native/rand.h +7 -7
- 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 +22 -22
- warp/native/sparse.cpp +8 -8
- warp/native/sparse.cu +14 -14
- warp/native/spatial.h +366 -17
- warp/native/svd.h +23 -8
- warp/native/temp_buffer.h +2 -2
- warp/native/tile.h +303 -70
- warp/native/tile_radix_sort.h +5 -1
- warp/native/tile_reduce.h +16 -25
- warp/native/tuple.h +2 -2
- warp/native/vec.h +385 -18
- 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 +337 -193
- warp/native/warp.h +227 -226
- warp/optim/linear.py +736 -271
- warp/render/imgui_manager.py +289 -0
- warp/render/render_opengl.py +137 -57
- warp/render/render_usd.py +0 -1
- warp/sim/collide.py +1 -2
- warp/sim/graph_coloring.py +2 -2
- warp/sim/integrator_vbd.py +10 -2
- warp/sparse.py +559 -176
- warp/tape.py +2 -0
- 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_marching_cubes.py +233 -12
- warp/tests/sim/test_cloth.py +89 -6
- warp/tests/sim/test_coloring.py +82 -7
- warp/tests/test_array.py +56 -5
- warp/tests/test_assert.py +53 -0
- warp/tests/test_atomic_cas.py +127 -114
- warp/tests/test_codegen.py +3 -2
- 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 +45 -2
- 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 +1 -1
- warp/tests/test_mat.py +1540 -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 +162 -34
- warp/tests/test_quat_assign_copy.py +145 -0
- warp/tests/test_reload.py +2 -1
- warp/tests/test_sparse.py +103 -0
- warp/tests/test_spatial.py +140 -34
- warp/tests/test_spatial_assign_copy.py +160 -0
- warp/tests/test_static.py +48 -0
- warp/tests/test_struct.py +43 -3
- warp/tests/test_tape.py +38 -0
- warp/tests/test_types.py +0 -20
- warp/tests/test_vec.py +216 -441
- warp/tests/test_vec_assign_copy.py +143 -0
- warp/tests/test_vec_constructors.py +325 -0
- warp/tests/tile/test_tile.py +206 -152
- 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 +179 -0
- warp/tests/tile/test_tile_mlp.py +1 -1
- warp/tests/tile/test_tile_reduce.py +100 -11
- warp/tests/tile/test_tile_shared_memory.py +16 -16
- warp/tests/tile/test_tile_sort.py +59 -55
- warp/tests/unittest_suites.py +16 -0
- warp/tests/walkthrough_debug.py +1 -1
- warp/thirdparty/unittest_parallel.py +108 -9
- warp/types.py +554 -264
- warp/utils.py +68 -86
- {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/METADATA +28 -65
- {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/RECORD +150 -138
- warp/native/marching.cpp +0 -19
- warp/native/marching.cu +0 -514
- warp/native/marching.h +0 -19
- {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/WHEEL +0 -0
- {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/licenses/LICENSE.md +0 -0
- {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/top_level.txt +0 -0
warp/types.py
CHANGED
|
@@ -176,11 +176,11 @@ def constant(x):
|
|
|
176
176
|
|
|
177
177
|
|
|
178
178
|
def float_to_half_bits(value):
|
|
179
|
-
return warp.context.runtime.core.
|
|
179
|
+
return warp.context.runtime.core.wp_float_to_half_bits(value)
|
|
180
180
|
|
|
181
181
|
|
|
182
182
|
def half_bits_to_float(value):
|
|
183
|
-
return warp.context.runtime.core.
|
|
183
|
+
return warp.context.runtime.core.wp_half_bits_to_float(value)
|
|
184
184
|
|
|
185
185
|
|
|
186
186
|
def safe_len(obj):
|
|
@@ -190,10 +190,80 @@ def safe_len(obj):
|
|
|
190
190
|
return -1
|
|
191
191
|
|
|
192
192
|
|
|
193
|
+
def flatten(value: Sequence) -> tuple[list, tuple[int, ...]]:
|
|
194
|
+
"""Flatten an arbitrarily-nested, rectangular iterable."""
|
|
195
|
+
arr = []
|
|
196
|
+
shape = []
|
|
197
|
+
|
|
198
|
+
depth = 0
|
|
199
|
+
stack = [(depth, value)]
|
|
200
|
+
|
|
201
|
+
while stack:
|
|
202
|
+
depth, elem = stack.pop(0)
|
|
203
|
+
|
|
204
|
+
if isinstance(elem, (str, bytes, bytearray, memoryview)):
|
|
205
|
+
raise TypeError(f"Got an invalid element of type `{type(elem).__name__}`")
|
|
206
|
+
|
|
207
|
+
try:
|
|
208
|
+
# If `elem` is a sequence, then it should be possible
|
|
209
|
+
# to add its elements to the stack for later processing.
|
|
210
|
+
stack.extend((depth + 1, x) for x in elem)
|
|
211
|
+
except TypeError:
|
|
212
|
+
# Since `elem` doesn't seem to be a sequence, we must have
|
|
213
|
+
# a leaf value that we need to add to our resulting array.
|
|
214
|
+
if depth != len(shape):
|
|
215
|
+
raise ValueError("Ragged array: scalar found before deepest level.") from None
|
|
216
|
+
|
|
217
|
+
arr.append(elem)
|
|
218
|
+
else:
|
|
219
|
+
dim = len(elem)
|
|
220
|
+
if depth == len(shape):
|
|
221
|
+
# First sequence seen at this depth, record its length.
|
|
222
|
+
shape.append(dim)
|
|
223
|
+
elif shape[depth] != dim:
|
|
224
|
+
# Later sequences must have the same length.
|
|
225
|
+
raise ValueError(f"Ragged array: expected length {shape[depth]} at depth {depth}, got {dim}.") from None
|
|
226
|
+
|
|
227
|
+
return (arr, tuple(shape))
|
|
228
|
+
|
|
229
|
+
|
|
193
230
|
# ----------------------
|
|
194
231
|
# built-in types
|
|
195
232
|
|
|
196
233
|
|
|
234
|
+
def _unary_op(self, op, t):
|
|
235
|
+
try:
|
|
236
|
+
return op(self)
|
|
237
|
+
except RuntimeError:
|
|
238
|
+
return t(*(op(a) for a in self))
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
def _binary_op(self, op, x, t, cw=True):
|
|
242
|
+
try:
|
|
243
|
+
return op(self, x)
|
|
244
|
+
except RuntimeError:
|
|
245
|
+
if is_scalar(x):
|
|
246
|
+
return t(*(op(a, x) for a in self))
|
|
247
|
+
|
|
248
|
+
if cw and types_equal(x, t):
|
|
249
|
+
return t(*(op(a, b) for a, b in zip(self, x)))
|
|
250
|
+
|
|
251
|
+
raise
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
def _rbinary_op(self, op, x, t, cw=True):
|
|
255
|
+
try:
|
|
256
|
+
return op(x, self)
|
|
257
|
+
except RuntimeError:
|
|
258
|
+
if is_scalar(x):
|
|
259
|
+
return t(*(op(x, a) for a in self))
|
|
260
|
+
|
|
261
|
+
if cw and types_equal(x, t):
|
|
262
|
+
return t(*(op(b, a) for a, b in zip(self, x)))
|
|
263
|
+
|
|
264
|
+
raise
|
|
265
|
+
|
|
266
|
+
|
|
197
267
|
def vector(length, dtype):
|
|
198
268
|
# canonicalize dtype
|
|
199
269
|
if dtype == int:
|
|
@@ -260,9 +330,10 @@ def vector(length, dtype):
|
|
|
260
330
|
return vec_t.scalar_export(super().__getitem__(key))
|
|
261
331
|
elif isinstance(key, slice):
|
|
262
332
|
if self._wp_scalar_type_ == float16:
|
|
263
|
-
|
|
333
|
+
values = tuple(vec_t.scalar_export(x) for x in super().__getitem__(key))
|
|
264
334
|
else:
|
|
265
|
-
|
|
335
|
+
values = super().__getitem__(key)
|
|
336
|
+
return vector(len(values), self._wp_scalar_type_)(*values)
|
|
266
337
|
else:
|
|
267
338
|
raise KeyError(f"Invalid key {key}, expected int or slice")
|
|
268
339
|
|
|
@@ -276,6 +347,13 @@ def vector(length, dtype):
|
|
|
276
347
|
f"but got `{type(value).__name__}` instead"
|
|
277
348
|
) from None
|
|
278
349
|
elif isinstance(key, slice):
|
|
350
|
+
if is_scalar(value):
|
|
351
|
+
indices = range(*key.indices(self._length_))
|
|
352
|
+
for idx in indices:
|
|
353
|
+
super().__setitem__(idx, vec_t.scalar_import(value))
|
|
354
|
+
|
|
355
|
+
return
|
|
356
|
+
|
|
279
357
|
try:
|
|
280
358
|
iter(value)
|
|
281
359
|
except TypeError:
|
|
@@ -325,40 +403,40 @@ def vector(length, dtype):
|
|
|
325
403
|
return super().__setattr__(name, value)
|
|
326
404
|
|
|
327
405
|
def __add__(self, y):
|
|
328
|
-
return warp.add
|
|
406
|
+
return _binary_op(self, warp.add, y, vec_t)
|
|
329
407
|
|
|
330
408
|
def __radd__(self, y):
|
|
331
|
-
return warp.add
|
|
409
|
+
return _rbinary_op(self, warp.add, y, vec_t)
|
|
332
410
|
|
|
333
411
|
def __sub__(self, y):
|
|
334
|
-
return warp.sub
|
|
412
|
+
return _binary_op(self, warp.sub, y, vec_t)
|
|
335
413
|
|
|
336
414
|
def __rsub__(self, y):
|
|
337
|
-
return warp.sub
|
|
415
|
+
return _rbinary_op(self, warp.sub, y, vec_t)
|
|
338
416
|
|
|
339
417
|
def __mul__(self, y):
|
|
340
|
-
return warp.mul
|
|
418
|
+
return _binary_op(self, warp.mul, y, vec_t, cw=False)
|
|
341
419
|
|
|
342
420
|
def __rmul__(self, x):
|
|
343
|
-
return warp.mul
|
|
421
|
+
return _rbinary_op(self, warp.mul, x, vec_t, cw=False)
|
|
344
422
|
|
|
345
423
|
def __truediv__(self, y):
|
|
346
|
-
return warp.div
|
|
424
|
+
return _binary_op(self, warp.div, y, vec_t, cw=False)
|
|
347
425
|
|
|
348
426
|
def __rtruediv__(self, x):
|
|
349
|
-
return warp.div
|
|
427
|
+
return _rbinary_op(self, warp.div, x, vec_t, cw=False)
|
|
350
428
|
|
|
351
429
|
def __mod__(self, x):
|
|
352
|
-
return warp.mod
|
|
430
|
+
return _binary_op(self, warp.mod, x, vec_t)
|
|
353
431
|
|
|
354
432
|
def __rmod__(self, x):
|
|
355
|
-
return warp.mod
|
|
433
|
+
return _rbinary_op(self, warp.mod, x, vec_t)
|
|
356
434
|
|
|
357
435
|
def __pos__(self):
|
|
358
|
-
return warp.pos
|
|
436
|
+
return _unary_op(self, warp.pos, vec_t)
|
|
359
437
|
|
|
360
438
|
def __neg__(self):
|
|
361
|
-
return warp.neg
|
|
439
|
+
return _unary_op(self, warp.neg, vec_t)
|
|
362
440
|
|
|
363
441
|
def __str__(self):
|
|
364
442
|
return f"[{', '.join(map(str, self))}]"
|
|
@@ -422,6 +500,7 @@ def matrix(shape, dtype):
|
|
|
422
500
|
_wp_constructor_ = "matrix"
|
|
423
501
|
|
|
424
502
|
_wp_row_type_ = vector(0 if shape[1] == Any else shape[1], dtype)
|
|
503
|
+
_wp_col_type_ = vector(0 if shape[0] == Any else shape[0], dtype)
|
|
425
504
|
|
|
426
505
|
# special handling for float16 type: in this case, data is stored
|
|
427
506
|
# as uint16 but it's actually half precision floating point
|
|
@@ -467,22 +546,22 @@ def matrix(shape, dtype):
|
|
|
467
546
|
return self._shape_[0]
|
|
468
547
|
|
|
469
548
|
def __add__(self, y):
|
|
470
|
-
return warp.add
|
|
549
|
+
return _binary_op(self, warp.add, y, mat_t)
|
|
471
550
|
|
|
472
551
|
def __radd__(self, y):
|
|
473
|
-
return warp.add
|
|
552
|
+
return _rbinary_op(self, warp.add, y, mat_t)
|
|
474
553
|
|
|
475
554
|
def __sub__(self, y):
|
|
476
|
-
return warp.sub
|
|
555
|
+
return _binary_op(self, warp.sub, y, mat_t)
|
|
477
556
|
|
|
478
557
|
def __rsub__(self, y):
|
|
479
|
-
return warp.sub
|
|
558
|
+
return _rbinary_op(self, warp.sub, y, mat_t)
|
|
480
559
|
|
|
481
560
|
def __mul__(self, y):
|
|
482
|
-
return warp.mul
|
|
561
|
+
return _binary_op(self, warp.mul, y, mat_t, cw=False)
|
|
483
562
|
|
|
484
563
|
def __rmul__(self, x):
|
|
485
|
-
return warp.mul
|
|
564
|
+
return _rbinary_op(self, warp.mul, x, mat_t, cw=False)
|
|
486
565
|
|
|
487
566
|
def __matmul__(self, y):
|
|
488
567
|
return warp.mul(self, y)
|
|
@@ -491,16 +570,22 @@ def matrix(shape, dtype):
|
|
|
491
570
|
return warp.mul(x, self)
|
|
492
571
|
|
|
493
572
|
def __truediv__(self, y):
|
|
494
|
-
return warp.div
|
|
573
|
+
return _binary_op(self, warp.div, y, mat_t, cw=False)
|
|
495
574
|
|
|
496
575
|
def __rtruediv__(self, x):
|
|
497
|
-
return warp.div
|
|
576
|
+
return _rbinary_op(self, warp.div, x, mat_t, cw=False)
|
|
577
|
+
|
|
578
|
+
def __mod__(self, x):
|
|
579
|
+
return _binary_op(self, warp.mod, x, mat_t)
|
|
580
|
+
|
|
581
|
+
def __rmod__(self, x):
|
|
582
|
+
return _rbinary_op(self, warp.mod, x, mat_t)
|
|
498
583
|
|
|
499
584
|
def __pos__(self):
|
|
500
|
-
return warp.pos
|
|
585
|
+
return _unary_op(self, warp.pos, mat_t)
|
|
501
586
|
|
|
502
587
|
def __neg__(self):
|
|
503
|
-
return warp.neg
|
|
588
|
+
return _unary_op(self, warp.neg, mat_t)
|
|
504
589
|
|
|
505
590
|
def __str__(self):
|
|
506
591
|
row_str = []
|
|
@@ -522,8 +607,10 @@ def matrix(shape, dtype):
|
|
|
522
607
|
return True
|
|
523
608
|
|
|
524
609
|
def get_row(self, r):
|
|
525
|
-
if r < 0 or r >= self._shape_[0]:
|
|
610
|
+
if r < -self._shape_[0] or r >= self._shape_[0]:
|
|
526
611
|
raise IndexError("Invalid row index")
|
|
612
|
+
if r < 0:
|
|
613
|
+
r += self._shape_[0]
|
|
527
614
|
row_start = r * self._shape_[1]
|
|
528
615
|
row_end = row_start + self._shape_[1]
|
|
529
616
|
row_data = super().__getitem__(slice(row_start, row_end))
|
|
@@ -532,9 +619,35 @@ def matrix(shape, dtype):
|
|
|
532
619
|
else:
|
|
533
620
|
return self._wp_row_type_(row_data)
|
|
534
621
|
|
|
622
|
+
def get_col(self, c):
|
|
623
|
+
if c < -self._shape_[1] or c >= self._shape_[1]:
|
|
624
|
+
raise IndexError("Invalid column index")
|
|
625
|
+
if c < 0:
|
|
626
|
+
c += self._shape_[1]
|
|
627
|
+
col_start = c
|
|
628
|
+
col_end = col_start + self._shape_[0] * self._shape_[1]
|
|
629
|
+
col_step = self._shape_[1]
|
|
630
|
+
col_data = super().__getitem__(slice(col_start, col_end, col_step))
|
|
631
|
+
if self._wp_scalar_type_ == float16:
|
|
632
|
+
return self._wp_col_type_(*[mat_t.scalar_export(x) for x in col_data])
|
|
633
|
+
else:
|
|
634
|
+
return self._wp_col_type_(col_data)
|
|
635
|
+
|
|
535
636
|
def set_row(self, r, v):
|
|
536
|
-
if r < 0 or r >= self._shape_[0]:
|
|
637
|
+
if r < -self._shape_[0] or r >= self._shape_[0]:
|
|
537
638
|
raise IndexError("Invalid row index")
|
|
639
|
+
if r < 0:
|
|
640
|
+
r += self._shape_[0]
|
|
641
|
+
|
|
642
|
+
row_start = r * self._shape_[1]
|
|
643
|
+
row_end = row_start + self._shape_[1]
|
|
644
|
+
|
|
645
|
+
if is_scalar(v):
|
|
646
|
+
for i in range(row_start, row_end):
|
|
647
|
+
super().__setitem__(i, mat_t.scalar_import(v))
|
|
648
|
+
|
|
649
|
+
return
|
|
650
|
+
|
|
538
651
|
try:
|
|
539
652
|
iter(v)
|
|
540
653
|
except TypeError:
|
|
@@ -542,8 +655,6 @@ def matrix(shape, dtype):
|
|
|
542
655
|
f"Expected to assign a slice from a sequence of values but got `{type(v).__name__}` instead"
|
|
543
656
|
) from None
|
|
544
657
|
|
|
545
|
-
row_start = r * self._shape_[1]
|
|
546
|
-
row_end = row_start + self._shape_[1]
|
|
547
658
|
if self._wp_scalar_type_ == float16:
|
|
548
659
|
converted = []
|
|
549
660
|
try:
|
|
@@ -558,17 +669,86 @@ def matrix(shape, dtype):
|
|
|
558
669
|
v = converted
|
|
559
670
|
super().__setitem__(slice(row_start, row_end), v)
|
|
560
671
|
|
|
672
|
+
def set_col(self, c, v):
|
|
673
|
+
if c < -self._shape_[1] or c >= self._shape_[1]:
|
|
674
|
+
raise IndexError("Invalid col index")
|
|
675
|
+
if c < 0:
|
|
676
|
+
c += self._shape_[1]
|
|
677
|
+
|
|
678
|
+
col_start = c
|
|
679
|
+
col_end = col_start + self._shape_[0] * self._shape_[1]
|
|
680
|
+
col_step = self._shape_[1]
|
|
681
|
+
|
|
682
|
+
if is_scalar(v):
|
|
683
|
+
for i in range(col_start, col_end, col_step):
|
|
684
|
+
super().__setitem__(i, mat_t.scalar_import(v))
|
|
685
|
+
|
|
686
|
+
return
|
|
687
|
+
|
|
688
|
+
try:
|
|
689
|
+
iter(v)
|
|
690
|
+
except TypeError:
|
|
691
|
+
raise TypeError(
|
|
692
|
+
f"Expected to assign a slice from a sequence of values but got `{type(v).__name__}` instead"
|
|
693
|
+
) from None
|
|
694
|
+
|
|
695
|
+
if self._wp_scalar_type_ == float16:
|
|
696
|
+
converted = []
|
|
697
|
+
try:
|
|
698
|
+
for x in v:
|
|
699
|
+
converted.append(mat_t.scalar_import(x))
|
|
700
|
+
except ctypes.ArgumentError:
|
|
701
|
+
raise TypeError(
|
|
702
|
+
f"Expected to assign a slice from a sequence of `float16` values "
|
|
703
|
+
f"but got `{type(x).__name__}` instead"
|
|
704
|
+
) from None
|
|
705
|
+
|
|
706
|
+
v = converted
|
|
707
|
+
super().__setitem__(slice(col_start, col_end, col_step), v)
|
|
708
|
+
|
|
561
709
|
def __getitem__(self, key):
|
|
562
710
|
if isinstance(key, Tuple):
|
|
563
711
|
# element indexing m[i,j]
|
|
564
712
|
if len(key) != 2:
|
|
565
713
|
raise KeyError(f"Invalid key, expected one or two indices, got {len(key)}")
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
|
|
714
|
+
|
|
715
|
+
# Count how many dimensions the output value will have.
|
|
716
|
+
ndim = sum(1 for x in key if isinstance(x, slice))
|
|
717
|
+
|
|
718
|
+
if ndim == 0:
|
|
719
|
+
row = key[0] + self._shape_[0] if key[0] < 0 else key[0]
|
|
720
|
+
col = key[1] + self._shape_[1] if key[1] < 0 else key[1]
|
|
721
|
+
return mat_t.scalar_export(super().__getitem__(row * self._shape_[1] + col))
|
|
722
|
+
|
|
723
|
+
if ndim == 1:
|
|
724
|
+
if isinstance(key[1], slice):
|
|
725
|
+
# Row vector.
|
|
726
|
+
cols = range(*key[1].indices(self._shape_[0]))
|
|
727
|
+
row_vec = self.get_row(key[0])
|
|
728
|
+
values = tuple(row_vec[x] for x in cols)
|
|
729
|
+
return vector(len(values), self._wp_scalar_type_)(*values)
|
|
730
|
+
else:
|
|
731
|
+
# Column vector.
|
|
732
|
+
rows = range(*key[0].indices(self._shape_[1]))
|
|
733
|
+
col_vec = self.get_col(key[1])
|
|
734
|
+
values = tuple(col_vec[x] for x in rows)
|
|
735
|
+
return vector(len(values), self._wp_scalar_type_)(*values)
|
|
736
|
+
|
|
737
|
+
assert ndim == 2
|
|
738
|
+
rows = range(*key[0].indices(self._shape_[1]))
|
|
739
|
+
cols = range(*key[1].indices(self._shape_[0]))
|
|
740
|
+
row_vecs = tuple(self.get_row(i) for i in rows)
|
|
741
|
+
values = tuple(x[j] for x in row_vecs for j in cols)
|
|
742
|
+
shape = (len(rows), len(cols))
|
|
743
|
+
return matrix(shape, self._wp_scalar_type_)(*values)
|
|
569
744
|
elif isinstance(key, int):
|
|
570
745
|
# row vector indexing m[r]
|
|
571
746
|
return self.get_row(key)
|
|
747
|
+
elif isinstance(key, slice):
|
|
748
|
+
indices = range(*key.indices(self._shape_[0]))
|
|
749
|
+
row_vecs = tuple(self.get_row(x) for x in indices)
|
|
750
|
+
shape = (len(row_vecs), self._shape_[1])
|
|
751
|
+
return matrix(shape, self._wp_scalar_type_)(*row_vecs)
|
|
572
752
|
else:
|
|
573
753
|
raise KeyError(f"Invalid key {key}, expected int or pair of ints")
|
|
574
754
|
|
|
@@ -577,20 +757,104 @@ def matrix(shape, dtype):
|
|
|
577
757
|
# element indexing m[i,j] = x
|
|
578
758
|
if len(key) != 2:
|
|
579
759
|
raise KeyError(f"Invalid key, expected one or two indices, got {len(key)}")
|
|
580
|
-
|
|
581
|
-
|
|
582
|
-
|
|
583
|
-
|
|
584
|
-
|
|
585
|
-
|
|
586
|
-
|
|
587
|
-
|
|
588
|
-
|
|
760
|
+
|
|
761
|
+
# Count how many dimensions the output value is expected to have.
|
|
762
|
+
ndim = sum(1 for x in key if isinstance(x, slice))
|
|
763
|
+
|
|
764
|
+
if ndim == 0:
|
|
765
|
+
try:
|
|
766
|
+
_, v_shape = flatten(value)
|
|
767
|
+
except TypeError:
|
|
768
|
+
raise TypeError(
|
|
769
|
+
f"Expected to assign a `{type_repr(self._wp_scalar_type_)}` value but got `{type(value).__name__}` instead"
|
|
770
|
+
) from None
|
|
771
|
+
|
|
772
|
+
if v_shape:
|
|
773
|
+
raise RuntimeError(
|
|
774
|
+
f"The provided value is expected to be a scalar but got an object of shape {v_shape} instead"
|
|
775
|
+
)
|
|
776
|
+
|
|
777
|
+
row = key[0] + self._shape_[0] if key[0] < 0 else key[0]
|
|
778
|
+
col = key[1] + self._shape_[1] if key[1] < 0 else key[1]
|
|
779
|
+
idx = row * self._shape_[1] + col
|
|
780
|
+
super().__setitem__(idx, mat_t.scalar_import(value))
|
|
781
|
+
return
|
|
782
|
+
|
|
783
|
+
if ndim == 1:
|
|
784
|
+
_, v_shape = flatten(value)
|
|
785
|
+
|
|
786
|
+
if v_shape and len(v_shape) != 1:
|
|
787
|
+
raise RuntimeError(
|
|
788
|
+
f"The provided value is expected to be a 1D vector but got an object of shape {v_shape} instead"
|
|
789
|
+
)
|
|
790
|
+
|
|
791
|
+
if isinstance(key[1], slice):
|
|
792
|
+
# Row vector.
|
|
793
|
+
cols = range(*key[1].indices(self._shape_[0]))
|
|
794
|
+
if v_shape and v_shape[0] != len(cols):
|
|
795
|
+
raise RuntimeError(
|
|
796
|
+
f"The length of the provided vector ({v_shape[0]}) isn't compatible with the given slice (expected {len(cols)})"
|
|
797
|
+
)
|
|
798
|
+
|
|
799
|
+
row = key[0] + self._shape_[0] if key[0] < 0 else key[0]
|
|
800
|
+
for i, col in enumerate(cols):
|
|
801
|
+
idx = row * self._shape_[1] + col
|
|
802
|
+
super().__setitem__(idx, mat_t.scalar_import(value[i] if v_shape else value))
|
|
803
|
+
|
|
804
|
+
return
|
|
805
|
+
else:
|
|
806
|
+
# Column vector.
|
|
807
|
+
rows = range(*key[0].indices(self._shape_[1]))
|
|
808
|
+
if v_shape and v_shape[0] != len(rows):
|
|
809
|
+
raise RuntimeError(
|
|
810
|
+
f"The length of the provided vector ({v_shape[0]}) isn't compatible with the given slice (expected {len(rows)})"
|
|
811
|
+
)
|
|
812
|
+
|
|
813
|
+
col = key[1] + self._shape_[1] if key[1] < 0 else key[1]
|
|
814
|
+
for i, row in enumerate(rows):
|
|
815
|
+
idx = row * self._shape_[1] + col
|
|
816
|
+
super().__setitem__(idx, mat_t.scalar_import(value[i] if v_shape else value))
|
|
817
|
+
|
|
818
|
+
return
|
|
819
|
+
|
|
820
|
+
assert ndim == 2
|
|
821
|
+
|
|
822
|
+
_, v_shape = flatten(value)
|
|
823
|
+
|
|
824
|
+
if v_shape and len(v_shape) != 2:
|
|
825
|
+
raise RuntimeError(
|
|
826
|
+
f"The provided value is expected to be a 2D matrix but got an object of shape {v_shape} instead"
|
|
827
|
+
)
|
|
828
|
+
|
|
829
|
+
rows = range(*key[0].indices(self._shape_[1]))
|
|
830
|
+
cols = range(*key[1].indices(self._shape_[0]))
|
|
831
|
+
|
|
832
|
+
if v_shape and v_shape != (len(rows), len(cols)):
|
|
833
|
+
raise RuntimeError(
|
|
834
|
+
f"The shape of the provided matrix ({v_shape}) isn't compatible with the given slice (expected ({len(rows)}, {len(cols)}))"
|
|
835
|
+
)
|
|
836
|
+
|
|
837
|
+
for i, row in enumerate(rows):
|
|
838
|
+
for j, col in enumerate(cols):
|
|
839
|
+
idx = row * self._shape_[1] + col
|
|
840
|
+
super().__setitem__(idx, mat_t.scalar_import(value[i, j] if v_shape else value))
|
|
589
841
|
elif isinstance(key, int):
|
|
590
842
|
# row vector indexing m[r] = v
|
|
591
843
|
return self.set_row(key, value)
|
|
592
844
|
elif isinstance(key, slice):
|
|
593
|
-
|
|
845
|
+
v_arr, v_shape = flatten(value)
|
|
846
|
+
indices = range(*key.indices(self._shape_[0]))
|
|
847
|
+
|
|
848
|
+
if v_shape and (len(v_shape) != 2 or v_shape[0] != len(indices) or v_shape[1] != self._shape_[1]):
|
|
849
|
+
raise RuntimeError(
|
|
850
|
+
f"The shape of the provided matrix ({v_shape}) isn't compatible with the given slice (expected ({len(indices)}, {self._shape_[1]}))"
|
|
851
|
+
)
|
|
852
|
+
|
|
853
|
+
for i, row in enumerate(indices):
|
|
854
|
+
offset = i * self._shape_[1]
|
|
855
|
+
self.set_row(
|
|
856
|
+
row, v_arr[offset : offset + self._shape_[1]] if v_shape else (value,) * self._shape_[1]
|
|
857
|
+
)
|
|
594
858
|
else:
|
|
595
859
|
raise KeyError(f"Invalid key {key}, expected int or pair of ints")
|
|
596
860
|
|
|
@@ -608,6 +872,60 @@ def matrix(shape, dtype):
|
|
|
608
872
|
return mat_t
|
|
609
873
|
|
|
610
874
|
|
|
875
|
+
def matrix_from_cols(*args: Sequence[Vector]):
|
|
876
|
+
if not all(type_is_vector(x) for x in args):
|
|
877
|
+
raise RuntimeError("all arguments are expected to be vectors")
|
|
878
|
+
|
|
879
|
+
length = args[0]._length_
|
|
880
|
+
if any(x._length_ != length for x in args):
|
|
881
|
+
raise RuntimeError("all vectors are expected to have the same length")
|
|
882
|
+
|
|
883
|
+
dtype = args[0]._wp_scalar_type_
|
|
884
|
+
if any(x._wp_scalar_type_ != dtype for x in args):
|
|
885
|
+
raise RuntimeError("all vectors are expected to have the same dtype")
|
|
886
|
+
|
|
887
|
+
row_count = length
|
|
888
|
+
col_count = len(args)
|
|
889
|
+
out = matrix(shape=(row_count, col_count), dtype=dtype)()
|
|
890
|
+
mat_t = type(out)
|
|
891
|
+
|
|
892
|
+
for col in range(col_count):
|
|
893
|
+
v = args[col]
|
|
894
|
+
for row in range(row_count):
|
|
895
|
+
idx = col_count * row + col
|
|
896
|
+
value = mat_t.scalar_import(v[row])
|
|
897
|
+
super(mat_t, out).__setitem__(idx, value)
|
|
898
|
+
|
|
899
|
+
return out
|
|
900
|
+
|
|
901
|
+
|
|
902
|
+
def matrix_from_rows(*args: Sequence[Vector]):
|
|
903
|
+
if not all(type_is_vector(x) for x in args):
|
|
904
|
+
raise RuntimeError("all arguments are expected to be vectors")
|
|
905
|
+
|
|
906
|
+
length = args[0]._length_
|
|
907
|
+
if any(x._length_ != length for x in args):
|
|
908
|
+
raise RuntimeError("all vectors are expected to have the same length")
|
|
909
|
+
|
|
910
|
+
dtype = args[0]._wp_scalar_type_
|
|
911
|
+
if any(x._wp_scalar_type_ != dtype for x in args):
|
|
912
|
+
raise RuntimeError("all vectors are expected to have the same dtype")
|
|
913
|
+
|
|
914
|
+
row_count = len(args)
|
|
915
|
+
col_count = length
|
|
916
|
+
out = matrix(shape=(row_count, col_count), dtype=dtype)()
|
|
917
|
+
mat_t = type(out)
|
|
918
|
+
|
|
919
|
+
for row in range(row_count):
|
|
920
|
+
v = args[row]
|
|
921
|
+
for col in range(col_count):
|
|
922
|
+
idx = col_count * row + col
|
|
923
|
+
value = mat_t.scalar_import(v[col])
|
|
924
|
+
super(mat_t, out).__setitem__(idx, value)
|
|
925
|
+
|
|
926
|
+
return out
|
|
927
|
+
|
|
928
|
+
|
|
611
929
|
class void:
|
|
612
930
|
def __init__(self):
|
|
613
931
|
pass
|
|
@@ -664,12 +982,22 @@ class scalar_base:
|
|
|
664
982
|
|
|
665
983
|
|
|
666
984
|
class float_base(scalar_base):
|
|
667
|
-
|
|
985
|
+
def __str__(self) -> str:
|
|
986
|
+
return str(self.value)
|
|
987
|
+
|
|
988
|
+
def __repr__(self) -> str:
|
|
989
|
+
return f"{type(self).__name__}({self!s})"
|
|
668
990
|
|
|
669
991
|
|
|
670
992
|
class int_base(scalar_base):
|
|
671
993
|
def __index__(self) -> int:
|
|
672
|
-
return int(self.value)
|
|
994
|
+
return int(self._type_(self.value).value)
|
|
995
|
+
|
|
996
|
+
def __str__(self) -> str:
|
|
997
|
+
return str(self._type_(self.value).value)
|
|
998
|
+
|
|
999
|
+
def __repr__(self) -> str:
|
|
1000
|
+
return f"{type(self).__name__}({self!s})"
|
|
673
1001
|
|
|
674
1002
|
|
|
675
1003
|
class bool:
|
|
@@ -688,6 +1016,12 @@ class bool:
|
|
|
688
1016
|
def __int__(self) -> int:
|
|
689
1017
|
return int(self.value != 0)
|
|
690
1018
|
|
|
1019
|
+
def __str__(self) -> str:
|
|
1020
|
+
return str(self.value != 0)
|
|
1021
|
+
|
|
1022
|
+
def __repr__(self) -> str:
|
|
1023
|
+
return f"{type(self).__name__}({self!s})"
|
|
1024
|
+
|
|
691
1025
|
|
|
692
1026
|
class float16(float_base):
|
|
693
1027
|
_length_ = 1
|
|
@@ -1308,6 +1642,49 @@ class launch_bounds_t(ctypes.Structure):
|
|
|
1308
1642
|
self.shape[i] = 1
|
|
1309
1643
|
|
|
1310
1644
|
|
|
1645
|
+
INT_WIDTH = ctypes.sizeof(ctypes.c_int) * 8
|
|
1646
|
+
SLICE_BEGIN = (1 << (INT_WIDTH - 1)) - 1
|
|
1647
|
+
SLICE_END = -(1 << (INT_WIDTH - 1))
|
|
1648
|
+
|
|
1649
|
+
|
|
1650
|
+
class slice_t:
|
|
1651
|
+
_wp_native_name_ = "slice_t"
|
|
1652
|
+
|
|
1653
|
+
def __init__(self, start, stop, step):
|
|
1654
|
+
self.start = start
|
|
1655
|
+
self.stop = stop
|
|
1656
|
+
self.step = step
|
|
1657
|
+
|
|
1658
|
+
def get_length(self, parent_length, wrap=False):
|
|
1659
|
+
if any(isinstance(x, warp.codegen.Var) for x in (self.start, self.stop, self.step)):
|
|
1660
|
+
raise RuntimeError("Vector slice indices must be constant values.")
|
|
1661
|
+
|
|
1662
|
+
if self.step == 0:
|
|
1663
|
+
raise RuntimeError(f"Vector slice step {self.step} is invalid.")
|
|
1664
|
+
|
|
1665
|
+
if self.start == SLICE_BEGIN:
|
|
1666
|
+
start = parent_length - 1 if self.step < 0 else 0
|
|
1667
|
+
else:
|
|
1668
|
+
start = min(max(self.start, -parent_length), parent_length)
|
|
1669
|
+
if wrap:
|
|
1670
|
+
start = start + parent_length if start < 0 else start
|
|
1671
|
+
|
|
1672
|
+
if self.stop == SLICE_END:
|
|
1673
|
+
stop = -1 if self.step < 0 else parent_length
|
|
1674
|
+
else:
|
|
1675
|
+
stop = min(max(self.stop, -parent_length), parent_length)
|
|
1676
|
+
if wrap:
|
|
1677
|
+
stop = stop + parent_length if stop < 0 else stop
|
|
1678
|
+
|
|
1679
|
+
if self.step > 0 and start < stop:
|
|
1680
|
+
return 1 + (stop - start - 1) // self.step
|
|
1681
|
+
|
|
1682
|
+
if self.step < 0 and start > stop:
|
|
1683
|
+
return 1 + (start - stop - 1) // (-self.step)
|
|
1684
|
+
|
|
1685
|
+
return 0
|
|
1686
|
+
|
|
1687
|
+
|
|
1311
1688
|
class shape_t(ctypes.Structure):
|
|
1312
1689
|
_fields_ = (("dims", ctypes.c_int32 * ARRAY_MAX_DIMS),)
|
|
1313
1690
|
|
|
@@ -1388,6 +1765,13 @@ class tuple_t:
|
|
|
1388
1765
|
self.values = values
|
|
1389
1766
|
|
|
1390
1767
|
|
|
1768
|
+
class pointer_t:
|
|
1769
|
+
"""Used during codegen to represent pointer types."""
|
|
1770
|
+
|
|
1771
|
+
def __init__(self, dtype):
|
|
1772
|
+
self.dtype = dtype
|
|
1773
|
+
|
|
1774
|
+
|
|
1391
1775
|
def type_ctype(dtype):
|
|
1392
1776
|
if dtype == float:
|
|
1393
1777
|
return ctypes.c_float
|
|
@@ -1525,10 +1909,10 @@ def scalar_short_name(t):
|
|
|
1525
1909
|
# converts any known type to a human readable string, good for error messages, reporting etc
|
|
1526
1910
|
def type_repr(t) -> str:
|
|
1527
1911
|
if is_array(t):
|
|
1528
|
-
if t.device is None:
|
|
1912
|
+
if hasattr(t, "device") and t.device is None:
|
|
1529
1913
|
# array is used as a type annotation - display ndim instead of shape
|
|
1530
|
-
return f"
|
|
1531
|
-
return f"
|
|
1914
|
+
return f"{type(t).__name__}(ndim={t.ndim}, dtype={type_repr(t.dtype)})"
|
|
1915
|
+
return f"{type(t).__name__}(shape={t.shape}, dtype={type_repr(t.dtype)})"
|
|
1532
1916
|
if is_tuple(t):
|
|
1533
1917
|
return f"tuple({', '.join(type_repr(x) for x in t.types)})"
|
|
1534
1918
|
if is_tile(t):
|
|
@@ -1584,6 +1968,11 @@ def type_is_float(t):
|
|
|
1584
1968
|
return t in float_types
|
|
1585
1969
|
|
|
1586
1970
|
|
|
1971
|
+
# returns True if the passed *type* is a scalar
|
|
1972
|
+
def type_is_scalar(t):
|
|
1973
|
+
return type_is_int(t) or type_is_float(t)
|
|
1974
|
+
|
|
1975
|
+
|
|
1587
1976
|
# returns True if the passed *type* is a vector
|
|
1588
1977
|
def type_is_vector(t):
|
|
1589
1978
|
return getattr(t, "_wp_generic_type_hint_", None) is Vector
|
|
@@ -1621,6 +2010,10 @@ def is_float(x: Any) -> builtins.bool:
|
|
|
1621
2010
|
return type_is_float(type(x))
|
|
1622
2011
|
|
|
1623
2012
|
|
|
2013
|
+
def is_scalar(x: Any) -> builtins.bool:
|
|
2014
|
+
return type_is_scalar(type(x))
|
|
2015
|
+
|
|
2016
|
+
|
|
1624
2017
|
def is_value(x: Any) -> builtins.bool:
|
|
1625
2018
|
return type_is_value(type(x))
|
|
1626
2019
|
|
|
@@ -1750,7 +2143,11 @@ def types_equal(a, b, match_generic=False):
|
|
|
1750
2143
|
|
|
1751
2144
|
return True
|
|
1752
2145
|
|
|
1753
|
-
if
|
|
2146
|
+
if (
|
|
2147
|
+
is_array(a)
|
|
2148
|
+
and (issubclass(type(a), type(b)) or issubclass(type(b), type(a)))
|
|
2149
|
+
and types_equal(a.dtype, b.dtype, match_generic=match_generic)
|
|
2150
|
+
):
|
|
1754
2151
|
return True
|
|
1755
2152
|
|
|
1756
2153
|
# match NewStructInstance and Struct dtype
|
|
@@ -1863,10 +2260,6 @@ class array(Array[DType]):
|
|
|
1863
2260
|
taking two arguments: pointer and size. If ``None``, then no function is called.
|
|
1864
2261
|
"""
|
|
1865
2262
|
|
|
1866
|
-
# member attributes available during code-gen (e.g.: d = array.shape[0])
|
|
1867
|
-
# (initialized when needed)
|
|
1868
|
-
_vars = None
|
|
1869
|
-
|
|
1870
2263
|
def __new__(cls, *args, **kwargs):
|
|
1871
2264
|
instance = super().__new__(cls)
|
|
1872
2265
|
instance.deleter = None
|
|
@@ -2483,7 +2876,7 @@ class array(Array[DType]):
|
|
|
2483
2876
|
|
|
2484
2877
|
# Performance note: avoid wrapping the external stream in a temporary Stream object
|
|
2485
2878
|
if stream != array_stream.cuda_stream:
|
|
2486
|
-
warp.context.runtime.core.
|
|
2879
|
+
warp.context.runtime.core.wp_cuda_stream_wait_stream(
|
|
2487
2880
|
stream, array_stream.cuda_stream, array_stream.cached_event.cuda_event
|
|
2488
2881
|
)
|
|
2489
2882
|
elif self.device.is_cpu:
|
|
@@ -2722,10 +3115,10 @@ class array(Array[DType]):
|
|
|
2722
3115
|
@property
|
|
2723
3116
|
def vars(self):
|
|
2724
3117
|
# member attributes available during code-gen (e.g.: d = array.shape[0])
|
|
2725
|
-
|
|
2726
|
-
|
|
2727
|
-
|
|
2728
|
-
|
|
3118
|
+
return {
|
|
3119
|
+
"shape": warp.codegen.Var("shape", shape_t),
|
|
3120
|
+
"ptr": warp.codegen.Var("data", pointer_t(self.dtype)),
|
|
3121
|
+
}
|
|
2729
3122
|
|
|
2730
3123
|
def mark_init(self):
|
|
2731
3124
|
"""Resets this array's read flag"""
|
|
@@ -2830,11 +3223,11 @@ class array(Array[DType]):
|
|
|
2830
3223
|
carr_ptr = ctypes.pointer(carr)
|
|
2831
3224
|
|
|
2832
3225
|
if self.device.is_cuda:
|
|
2833
|
-
warp.context.runtime.core.
|
|
3226
|
+
warp.context.runtime.core.wp_array_fill_device(
|
|
2834
3227
|
self.device.context, carr_ptr, ARRAY_TYPE_REGULAR, cvalue_ptr, cvalue_size
|
|
2835
3228
|
)
|
|
2836
3229
|
else:
|
|
2837
|
-
warp.context.runtime.core.
|
|
3230
|
+
warp.context.runtime.core.wp_array_fill_host(carr_ptr, ARRAY_TYPE_REGULAR, cvalue_ptr, cvalue_size)
|
|
2838
3231
|
|
|
2839
3232
|
self.mark_init()
|
|
2840
3233
|
|
|
@@ -3138,7 +3531,7 @@ class array(Array[DType]):
|
|
|
3138
3531
|
# Allocate a buffer for the data (64-element char array)
|
|
3139
3532
|
ipc_handle_buffer = (ctypes.c_char * 64)()
|
|
3140
3533
|
|
|
3141
|
-
warp.context.runtime.core.
|
|
3534
|
+
warp.context.runtime.core.wp_cuda_ipc_get_mem_handle(self.ptr, ipc_handle_buffer)
|
|
3142
3535
|
|
|
3143
3536
|
return ipc_handle_buffer.raw
|
|
3144
3537
|
|
|
@@ -3191,7 +3584,7 @@ def from_ptr(ptr, length, dtype=None, shape=None, device=None):
|
|
|
3191
3584
|
|
|
3192
3585
|
|
|
3193
3586
|
def _close_cuda_ipc_handle(ptr, size):
|
|
3194
|
-
warp.context.runtime.core.
|
|
3587
|
+
warp.context.runtime.core.wp_cuda_ipc_close_mem_handle(ptr)
|
|
3195
3588
|
|
|
3196
3589
|
|
|
3197
3590
|
def from_ipc_handle(
|
|
@@ -3230,11 +3623,59 @@ def from_ipc_handle(
|
|
|
3230
3623
|
if device.is_ipc_supported is False:
|
|
3231
3624
|
raise RuntimeError(f"IPC is not supported on device {device}.")
|
|
3232
3625
|
|
|
3233
|
-
ptr = warp.context.runtime.core.
|
|
3626
|
+
ptr = warp.context.runtime.core.wp_cuda_ipc_open_mem_handle(device.context, handle)
|
|
3234
3627
|
|
|
3235
3628
|
return array(ptr=ptr, dtype=dtype, shape=shape, strides=strides, device=device, deleter=_close_cuda_ipc_handle)
|
|
3236
3629
|
|
|
3237
3630
|
|
|
3631
|
+
class fixedarray(array):
|
|
3632
|
+
"""A fixed-size, stack allocated, array containing values of the same type.
|
|
3633
|
+
|
|
3634
|
+
Only used during codegen, and for type hints, but otherwise not intended to be used
|
|
3635
|
+
at the Python scope.
|
|
3636
|
+
|
|
3637
|
+
Attributes:
|
|
3638
|
+
dtype (DType): The data type of the array.
|
|
3639
|
+
shape (tuple[int]): Dimensions of the array.
|
|
3640
|
+
"""
|
|
3641
|
+
|
|
3642
|
+
def __init__(
|
|
3643
|
+
self,
|
|
3644
|
+
dtype: Any = Any,
|
|
3645
|
+
shape: int | tuple[int, ...] | list[int] | None = None,
|
|
3646
|
+
):
|
|
3647
|
+
# canonicalize dtype
|
|
3648
|
+
if dtype == int:
|
|
3649
|
+
dtype = int32
|
|
3650
|
+
elif dtype == float:
|
|
3651
|
+
dtype = float32
|
|
3652
|
+
elif dtype == builtins.bool:
|
|
3653
|
+
dtype = bool
|
|
3654
|
+
|
|
3655
|
+
if shape is None:
|
|
3656
|
+
self.dtype = dtype
|
|
3657
|
+
self.ndim = 1
|
|
3658
|
+
self.size = 0
|
|
3659
|
+
self.shape = (0,)
|
|
3660
|
+
self.strides = (0,)
|
|
3661
|
+
else:
|
|
3662
|
+
if isinstance(shape, int):
|
|
3663
|
+
shape = (shape,)
|
|
3664
|
+
|
|
3665
|
+
check_array_shape(shape)
|
|
3666
|
+
|
|
3667
|
+
self.dtype = dtype
|
|
3668
|
+
self.ndim = len(shape)
|
|
3669
|
+
self.size = math.prod(shape)
|
|
3670
|
+
self.shape = shape
|
|
3671
|
+
self.strides = strides_from_shape(shape, dtype)
|
|
3672
|
+
|
|
3673
|
+
@property
|
|
3674
|
+
def vars(self):
|
|
3675
|
+
# member attributes available during code-gen (e.g.: d = array.shape[0])
|
|
3676
|
+
return {"shape": warp.codegen.Var("shape", shape_t)}
|
|
3677
|
+
|
|
3678
|
+
|
|
3238
3679
|
# A base class for non-contiguous arrays, providing the implementation of common methods like
|
|
3239
3680
|
# contiguous(), to(), numpy(), list(), assign(), zero_(), and fill_().
|
|
3240
3681
|
class noncontiguous_array_base(Array[T]):
|
|
@@ -3315,11 +3756,11 @@ class noncontiguous_array_base(Array[T]):
|
|
|
3315
3756
|
ctype_ptr = ctypes.pointer(ctype)
|
|
3316
3757
|
|
|
3317
3758
|
if self.device.is_cuda:
|
|
3318
|
-
warp.context.runtime.core.
|
|
3759
|
+
warp.context.runtime.core.wp_array_fill_device(
|
|
3319
3760
|
self.device.context, ctype_ptr, self.type_id, cvalue_ptr, cvalue_size
|
|
3320
3761
|
)
|
|
3321
3762
|
else:
|
|
3322
|
-
warp.context.runtime.core.
|
|
3763
|
+
warp.context.runtime.core.wp_array_fill_host(ctype_ptr, self.type_id, cvalue_ptr, cvalue_size)
|
|
3323
3764
|
|
|
3324
3765
|
|
|
3325
3766
|
# helper to check index array properties
|
|
@@ -3462,7 +3903,7 @@ def indexedarray4d(*args, **kwargs):
|
|
|
3462
3903
|
|
|
3463
3904
|
from warp.fabric import fabricarray, indexedfabricarray # noqa: E402
|
|
3464
3905
|
|
|
3465
|
-
array_types = (array, indexedarray, fabricarray, indexedfabricarray)
|
|
3906
|
+
array_types = (array, indexedarray, fabricarray, indexedfabricarray, fixedarray)
|
|
3466
3907
|
|
|
3467
3908
|
|
|
3468
3909
|
def array_type_id(a):
|
|
@@ -3663,11 +4104,11 @@ class Bvh:
|
|
|
3663
4104
|
)
|
|
3664
4105
|
constructor = "sah"
|
|
3665
4106
|
|
|
3666
|
-
self.id = self.runtime.core.
|
|
4107
|
+
self.id = self.runtime.core.wp_bvh_create_host(
|
|
3667
4108
|
get_data(lowers), get_data(uppers), len(lowers), bvh_constructor_values[constructor]
|
|
3668
4109
|
)
|
|
3669
4110
|
else:
|
|
3670
|
-
self.id = self.runtime.core.
|
|
4111
|
+
self.id = self.runtime.core.wp_bvh_create_device(
|
|
3671
4112
|
self.device.context,
|
|
3672
4113
|
get_data(lowers),
|
|
3673
4114
|
get_data(uppers),
|
|
@@ -3680,11 +4121,11 @@ class Bvh:
|
|
|
3680
4121
|
return
|
|
3681
4122
|
|
|
3682
4123
|
if self.device.is_cpu:
|
|
3683
|
-
self.runtime.core.
|
|
4124
|
+
self.runtime.core.wp_bvh_destroy_host(self.id)
|
|
3684
4125
|
else:
|
|
3685
4126
|
# use CUDA context guard to avoid side effects during garbage collection
|
|
3686
4127
|
with self.device.context_guard:
|
|
3687
|
-
self.runtime.core.
|
|
4128
|
+
self.runtime.core.wp_bvh_destroy_device(self.id)
|
|
3688
4129
|
|
|
3689
4130
|
def refit(self):
|
|
3690
4131
|
"""Refit the BVH.
|
|
@@ -3693,9 +4134,9 @@ class Bvh:
|
|
|
3693
4134
|
"""
|
|
3694
4135
|
|
|
3695
4136
|
if self.device.is_cpu:
|
|
3696
|
-
self.runtime.core.
|
|
4137
|
+
self.runtime.core.wp_bvh_refit_host(self.id)
|
|
3697
4138
|
else:
|
|
3698
|
-
self.runtime.core.
|
|
4139
|
+
self.runtime.core.wp_bvh_refit_device(self.id)
|
|
3699
4140
|
self.runtime.verify_cuda_device(self.device)
|
|
3700
4141
|
|
|
3701
4142
|
|
|
@@ -3777,7 +4218,7 @@ class Mesh:
|
|
|
3777
4218
|
)
|
|
3778
4219
|
bvh_constructor = "sah"
|
|
3779
4220
|
|
|
3780
|
-
self.id = self.runtime.core.
|
|
4221
|
+
self.id = self.runtime.core.wp_mesh_create_host(
|
|
3781
4222
|
points.__ctype__(),
|
|
3782
4223
|
velocities.__ctype__() if velocities else array().__ctype__(),
|
|
3783
4224
|
indices.__ctype__(),
|
|
@@ -3787,7 +4228,7 @@ class Mesh:
|
|
|
3787
4228
|
bvh_constructor_values[bvh_constructor],
|
|
3788
4229
|
)
|
|
3789
4230
|
else:
|
|
3790
|
-
self.id = self.runtime.core.
|
|
4231
|
+
self.id = self.runtime.core.wp_mesh_create_device(
|
|
3791
4232
|
self.device.context,
|
|
3792
4233
|
points.__ctype__(),
|
|
3793
4234
|
velocities.__ctype__() if velocities else array().__ctype__(),
|
|
@@ -3803,11 +4244,11 @@ class Mesh:
|
|
|
3803
4244
|
return
|
|
3804
4245
|
|
|
3805
4246
|
if self.device.is_cpu:
|
|
3806
|
-
self.runtime.core.
|
|
4247
|
+
self.runtime.core.wp_mesh_destroy_host(self.id)
|
|
3807
4248
|
else:
|
|
3808
4249
|
# use CUDA context guard to avoid side effects during garbage collection
|
|
3809
4250
|
with self.device.context_guard:
|
|
3810
|
-
self.runtime.core.
|
|
4251
|
+
self.runtime.core.wp_mesh_destroy_device(self.id)
|
|
3811
4252
|
|
|
3812
4253
|
def refit(self):
|
|
3813
4254
|
"""Refit the BVH to points.
|
|
@@ -3816,9 +4257,9 @@ class Mesh:
|
|
|
3816
4257
|
"""
|
|
3817
4258
|
|
|
3818
4259
|
if self.device.is_cpu:
|
|
3819
|
-
self.runtime.core.
|
|
4260
|
+
self.runtime.core.wp_mesh_refit_host(self.id)
|
|
3820
4261
|
else:
|
|
3821
|
-
self.runtime.core.
|
|
4262
|
+
self.runtime.core.wp_mesh_refit_device(self.id)
|
|
3822
4263
|
self.runtime.verify_cuda_device(self.device)
|
|
3823
4264
|
|
|
3824
4265
|
@property
|
|
@@ -3848,9 +4289,9 @@ class Mesh:
|
|
|
3848
4289
|
|
|
3849
4290
|
self._points = points_new
|
|
3850
4291
|
if self.device.is_cpu:
|
|
3851
|
-
self.runtime.core.
|
|
4292
|
+
self.runtime.core.wp_mesh_set_points_host(self.id, points_new.__ctype__())
|
|
3852
4293
|
else:
|
|
3853
|
-
self.runtime.core.
|
|
4294
|
+
self.runtime.core.wp_mesh_set_points_device(self.id, points_new.__ctype__())
|
|
3854
4295
|
self.runtime.verify_cuda_device(self.device)
|
|
3855
4296
|
|
|
3856
4297
|
@property
|
|
@@ -3878,9 +4319,9 @@ class Mesh:
|
|
|
3878
4319
|
|
|
3879
4320
|
self._velocities = velocities_new
|
|
3880
4321
|
if self.device.is_cpu:
|
|
3881
|
-
self.runtime.core.
|
|
4322
|
+
self.runtime.core.wp_mesh_set_velocities_host(self.id, velocities_new.__ctype__())
|
|
3882
4323
|
else:
|
|
3883
|
-
self.runtime.core.
|
|
4324
|
+
self.runtime.core.wp_mesh_set_velocities_device(self.id, velocities_new.__ctype__())
|
|
3884
4325
|
self.runtime.verify_cuda_device(self.device)
|
|
3885
4326
|
|
|
3886
4327
|
|
|
@@ -3912,11 +4353,11 @@ class Volume:
|
|
|
3912
4353
|
|
|
3913
4354
|
owner = False
|
|
3914
4355
|
if self.device.is_cpu:
|
|
3915
|
-
self.id = self.runtime.core.
|
|
4356
|
+
self.id = self.runtime.core.wp_volume_create_host(
|
|
3916
4357
|
ctypes.cast(data.ptr, ctypes.c_void_p), data.size, copy, owner
|
|
3917
4358
|
)
|
|
3918
4359
|
else:
|
|
3919
|
-
self.id = self.runtime.core.
|
|
4360
|
+
self.id = self.runtime.core.wp_volume_create_device(
|
|
3920
4361
|
self.device.context, ctypes.cast(data.ptr, ctypes.c_void_p), data.size, copy, owner
|
|
3921
4362
|
)
|
|
3922
4363
|
|
|
@@ -3928,18 +4369,18 @@ class Volume:
|
|
|
3928
4369
|
return
|
|
3929
4370
|
|
|
3930
4371
|
if self.device.is_cpu:
|
|
3931
|
-
self.runtime.core.
|
|
4372
|
+
self.runtime.core.wp_volume_destroy_host(self.id)
|
|
3932
4373
|
else:
|
|
3933
4374
|
# use CUDA context guard to avoid side effects during garbage collection
|
|
3934
4375
|
with self.device.context_guard:
|
|
3935
|
-
self.runtime.core.
|
|
4376
|
+
self.runtime.core.wp_volume_destroy_device(self.id)
|
|
3936
4377
|
|
|
3937
4378
|
def array(self) -> array:
|
|
3938
4379
|
"""Return the raw memory buffer of the :class:`Volume` as an array."""
|
|
3939
4380
|
|
|
3940
4381
|
buf = ctypes.c_void_p(0)
|
|
3941
4382
|
size = ctypes.c_uint64(0)
|
|
3942
|
-
self.runtime.core.
|
|
4383
|
+
self.runtime.core.wp_volume_get_buffer_info(self.id, ctypes.byref(buf), ctypes.byref(size))
|
|
3943
4384
|
return array(ptr=buf.value, dtype=uint8, shape=size.value, device=self.device)
|
|
3944
4385
|
|
|
3945
4386
|
def get_tile_count(self) -> int:
|
|
@@ -3949,7 +4390,9 @@ class Volume:
|
|
|
3949
4390
|
ctypes.c_uint64(0),
|
|
3950
4391
|
ctypes.c_uint32(0),
|
|
3951
4392
|
)
|
|
3952
|
-
self.runtime.core.
|
|
4393
|
+
self.runtime.core.wp_volume_get_tile_and_voxel_count(
|
|
4394
|
+
self.id, ctypes.byref(tile_count), ctypes.byref(voxel_count)
|
|
4395
|
+
)
|
|
3953
4396
|
return tile_count.value
|
|
3954
4397
|
|
|
3955
4398
|
def get_tiles(self, out: array | None = None) -> array:
|
|
@@ -3976,9 +4419,9 @@ class Volume:
|
|
|
3976
4419
|
)
|
|
3977
4420
|
|
|
3978
4421
|
if self.device.is_cpu:
|
|
3979
|
-
self.runtime.core.
|
|
4422
|
+
self.runtime.core.wp_volume_get_tiles_host(self.id, out.ptr)
|
|
3980
4423
|
else:
|
|
3981
|
-
self.runtime.core.
|
|
4424
|
+
self.runtime.core.wp_volume_get_tiles_device(self.id, out.ptr)
|
|
3982
4425
|
|
|
3983
4426
|
return out
|
|
3984
4427
|
|
|
@@ -3989,7 +4432,9 @@ class Volume:
|
|
|
3989
4432
|
ctypes.c_uint64(0),
|
|
3990
4433
|
ctypes.c_uint32(0),
|
|
3991
4434
|
)
|
|
3992
|
-
self.runtime.core.
|
|
4435
|
+
self.runtime.core.wp_volume_get_tile_and_voxel_count(
|
|
4436
|
+
self.id, ctypes.byref(tile_count), ctypes.byref(voxel_count)
|
|
4437
|
+
)
|
|
3993
4438
|
return voxel_count.value
|
|
3994
4439
|
|
|
3995
4440
|
def get_voxels(self, out: array | None = None) -> array:
|
|
@@ -4015,9 +4460,9 @@ class Volume:
|
|
|
4015
4460
|
)
|
|
4016
4461
|
|
|
4017
4462
|
if self.device.is_cpu:
|
|
4018
|
-
self.runtime.core.
|
|
4463
|
+
self.runtime.core.wp_volume_get_voxels_host(self.id, out.ptr)
|
|
4019
4464
|
else:
|
|
4020
|
-
self.runtime.core.
|
|
4465
|
+
self.runtime.core.wp_volume_get_voxels_device(self.id, out.ptr)
|
|
4021
4466
|
|
|
4022
4467
|
return out
|
|
4023
4468
|
|
|
@@ -4028,7 +4473,7 @@ class Volume:
|
|
|
4028
4473
|
raise RuntimeError("Invalid Volume")
|
|
4029
4474
|
|
|
4030
4475
|
dx, dy, dz = ctypes.c_float(0), ctypes.c_float(0), ctypes.c_float(0)
|
|
4031
|
-
self.runtime.core.
|
|
4476
|
+
self.runtime.core.wp_volume_get_voxel_size(self.id, ctypes.byref(dx), ctypes.byref(dy), ctypes.byref(dz))
|
|
4032
4477
|
return (dx.value, dy.value, dz.value)
|
|
4033
4478
|
|
|
4034
4479
|
class GridInfo(NamedTuple):
|
|
@@ -4061,7 +4506,7 @@ class Volume:
|
|
|
4061
4506
|
transform_buffer = (ctypes.c_float * 9)()
|
|
4062
4507
|
type_str_buffer = (ctypes.c_char * 16)()
|
|
4063
4508
|
|
|
4064
|
-
name = self.runtime.core.
|
|
4509
|
+
name = self.runtime.core.wp_volume_get_grid_info(
|
|
4065
4510
|
self.id,
|
|
4066
4511
|
ctypes.byref(grid_size),
|
|
4067
4512
|
ctypes.byref(grid_index),
|
|
@@ -4124,7 +4569,7 @@ class Volume:
|
|
|
4124
4569
|
def get_feature_array_count(self) -> int:
|
|
4125
4570
|
"""Return the number of supplemental data arrays stored alongside the grid"""
|
|
4126
4571
|
|
|
4127
|
-
return self.runtime.core.
|
|
4572
|
+
return self.runtime.core.wp_volume_get_blind_data_count(self.id)
|
|
4128
4573
|
|
|
4129
4574
|
class FeatureArrayInfo(NamedTuple):
|
|
4130
4575
|
"""Metadata for a supplemental data array"""
|
|
@@ -4149,7 +4594,7 @@ class Volume:
|
|
|
4149
4594
|
value_size = ctypes.c_uint32(0)
|
|
4150
4595
|
type_str_buffer = (ctypes.c_char * 16)()
|
|
4151
4596
|
|
|
4152
|
-
name = self.runtime.core.
|
|
4597
|
+
name = self.runtime.core.wp_volume_get_blind_data_info(
|
|
4153
4598
|
self.id,
|
|
4154
4599
|
feature_index,
|
|
4155
4600
|
ctypes.byref(buf),
|
|
@@ -4425,7 +4870,7 @@ class Volume:
|
|
|
4425
4870
|
# (to allow this we would need to ref-count the volume descriptor)
|
|
4426
4871
|
existing_buf = ctypes.c_void_p(0)
|
|
4427
4872
|
existing_size = ctypes.c_uint64(0)
|
|
4428
|
-
warp.context.runtime.core.
|
|
4873
|
+
warp.context.runtime.core.wp_volume_get_buffer_info(
|
|
4429
4874
|
grid_ptr, ctypes.byref(existing_buf), ctypes.byref(existing_size)
|
|
4430
4875
|
)
|
|
4431
4876
|
|
|
@@ -4678,7 +5123,7 @@ class Volume:
|
|
|
4678
5123
|
transform_buf, translation_buf = Volume._fill_transform_buffers(voxel_size, translation, transform)
|
|
4679
5124
|
|
|
4680
5125
|
if bg_value is None:
|
|
4681
|
-
volume.id = volume.runtime.core.
|
|
5126
|
+
volume.id = volume.runtime.core.wp_volume_index_from_tiles_device(
|
|
4682
5127
|
volume.device.context,
|
|
4683
5128
|
ctypes.c_void_p(tile_points.ptr),
|
|
4684
5129
|
tile_points.shape[0],
|
|
@@ -4717,7 +5162,7 @@ class Volume:
|
|
|
4717
5162
|
cvalue_size = ctypes.sizeof(cvalue)
|
|
4718
5163
|
cvalue_type = nvdb_type.encode("ascii")
|
|
4719
5164
|
|
|
4720
|
-
volume.id = volume.runtime.core.
|
|
5165
|
+
volume.id = volume.runtime.core.wp_volume_from_tiles_device(
|
|
4721
5166
|
volume.device.context,
|
|
4722
5167
|
ctypes.c_void_p(tile_points.ptr),
|
|
4723
5168
|
tile_points.shape[0],
|
|
@@ -4779,7 +5224,7 @@ class Volume:
|
|
|
4779
5224
|
|
|
4780
5225
|
transform_buf, translation_buf = Volume._fill_transform_buffers(voxel_size, translation, transform)
|
|
4781
5226
|
|
|
4782
|
-
volume.id = volume.runtime.core.
|
|
5227
|
+
volume.id = volume.runtime.core.wp_volume_from_active_voxels_device(
|
|
4783
5228
|
volume.device.context,
|
|
4784
5229
|
ctypes.c_void_p(voxel_points.ptr),
|
|
4785
5230
|
voxel_points.shape[0],
|
|
@@ -5019,9 +5464,9 @@ class HashGrid:
|
|
|
5019
5464
|
self.device = self.runtime.get_device(device)
|
|
5020
5465
|
|
|
5021
5466
|
if self.device.is_cpu:
|
|
5022
|
-
self.id = self.runtime.core.
|
|
5467
|
+
self.id = self.runtime.core.wp_hash_grid_create_host(dim_x, dim_y, dim_z)
|
|
5023
5468
|
else:
|
|
5024
|
-
self.id = self.runtime.core.
|
|
5469
|
+
self.id = self.runtime.core.wp_hash_grid_create_device(self.device.context, dim_x, dim_y, dim_z)
|
|
5025
5470
|
|
|
5026
5471
|
# indicates whether the grid data has been reserved for use by a kernel
|
|
5027
5472
|
self.reserved = False
|
|
@@ -5046,16 +5491,16 @@ class HashGrid:
|
|
|
5046
5491
|
points = points.contiguous().flatten()
|
|
5047
5492
|
|
|
5048
5493
|
if self.device.is_cpu:
|
|
5049
|
-
self.runtime.core.
|
|
5494
|
+
self.runtime.core.wp_hash_grid_update_host(self.id, radius, ctypes.byref(points.__ctype__()))
|
|
5050
5495
|
else:
|
|
5051
|
-
self.runtime.core.
|
|
5496
|
+
self.runtime.core.wp_hash_grid_update_device(self.id, radius, ctypes.byref(points.__ctype__()))
|
|
5052
5497
|
self.reserved = True
|
|
5053
5498
|
|
|
5054
5499
|
def reserve(self, num_points):
|
|
5055
5500
|
if self.device.is_cpu:
|
|
5056
|
-
self.runtime.core.
|
|
5501
|
+
self.runtime.core.wp_hash_grid_reserve_host(self.id, num_points)
|
|
5057
5502
|
else:
|
|
5058
|
-
self.runtime.core.
|
|
5503
|
+
self.runtime.core.wp_hash_grid_reserve_device(self.id, num_points)
|
|
5059
5504
|
self.reserved = True
|
|
5060
5505
|
|
|
5061
5506
|
def __del__(self):
|
|
@@ -5063,166 +5508,11 @@ class HashGrid:
|
|
|
5063
5508
|
return
|
|
5064
5509
|
|
|
5065
5510
|
if self.device.is_cpu:
|
|
5066
|
-
self.runtime.core.
|
|
5511
|
+
self.runtime.core.wp_hash_grid_destroy_host(self.id)
|
|
5067
5512
|
else:
|
|
5068
5513
|
# use CUDA context guard to avoid side effects during garbage collection
|
|
5069
5514
|
with self.device.context_guard:
|
|
5070
|
-
self.runtime.core.
|
|
5071
|
-
|
|
5072
|
-
|
|
5073
|
-
class MarchingCubes:
|
|
5074
|
-
def __new__(cls, *args, **kwargs):
|
|
5075
|
-
instance = super().__new__(cls)
|
|
5076
|
-
instance.id = None
|
|
5077
|
-
return instance
|
|
5078
|
-
|
|
5079
|
-
def __init__(self, nx: int, ny: int, nz: int, max_verts: int, max_tris: int, device=None):
|
|
5080
|
-
"""CUDA-based Marching Cubes algorithm to extract a 2D surface mesh from a 3D volume.
|
|
5081
|
-
|
|
5082
|
-
Attributes:
|
|
5083
|
-
id: Unique identifier for this object.
|
|
5084
|
-
verts (:class:`warp.array`): Array of vertex positions of type :class:`warp.vec3f`
|
|
5085
|
-
for the output surface mesh.
|
|
5086
|
-
This is populated after running :func:`surface`.
|
|
5087
|
-
indices (:class:`warp.array`): Array containing indices of type :class:`warp.int32`
|
|
5088
|
-
defining triangles for the output surface mesh.
|
|
5089
|
-
This is populated after running :func:`surface`.
|
|
5090
|
-
|
|
5091
|
-
Each set of three consecutive integers in the array represents a single triangle,
|
|
5092
|
-
in which each integer is an index referring to a vertex in the :attr:`verts` array.
|
|
5093
|
-
|
|
5094
|
-
Args:
|
|
5095
|
-
nx: Number of cubes in the x-direction.
|
|
5096
|
-
ny: Number of cubes in the y-direction.
|
|
5097
|
-
nz: Number of cubes in the z-direction.
|
|
5098
|
-
max_verts: Maximum expected number of vertices (used for array preallocation).
|
|
5099
|
-
max_tris: Maximum expected number of triangles (used for array preallocation).
|
|
5100
|
-
device (Devicelike): CUDA device on which to run marching cubes and allocate memory.
|
|
5101
|
-
|
|
5102
|
-
Raises:
|
|
5103
|
-
RuntimeError: ``device`` not a CUDA device.
|
|
5104
|
-
|
|
5105
|
-
.. note::
|
|
5106
|
-
The shape of the marching cubes should match the shape of the scalar field being surfaced.
|
|
5107
|
-
|
|
5108
|
-
"""
|
|
5109
|
-
|
|
5110
|
-
self.id = 0
|
|
5111
|
-
|
|
5112
|
-
self.runtime = warp.context.runtime
|
|
5113
|
-
|
|
5114
|
-
self.device = self.runtime.get_device(device)
|
|
5115
|
-
|
|
5116
|
-
if not self.device.is_cuda:
|
|
5117
|
-
raise RuntimeError("Only CUDA devices are supported for marching cubes")
|
|
5118
|
-
|
|
5119
|
-
self.nx = nx
|
|
5120
|
-
self.ny = ny
|
|
5121
|
-
self.nz = nz
|
|
5122
|
-
|
|
5123
|
-
self.max_verts = max_verts
|
|
5124
|
-
self.max_tris = max_tris
|
|
5125
|
-
|
|
5126
|
-
# bindings to warp.so
|
|
5127
|
-
self.alloc = self.runtime.core.marching_cubes_create_device
|
|
5128
|
-
self.alloc.argtypes = [ctypes.c_void_p]
|
|
5129
|
-
self.alloc.restype = ctypes.c_uint64
|
|
5130
|
-
self.free = self.runtime.core.marching_cubes_destroy_device
|
|
5131
|
-
|
|
5132
|
-
from warp.context import zeros
|
|
5133
|
-
|
|
5134
|
-
self.verts = zeros(max_verts, dtype=vec3, device=self.device)
|
|
5135
|
-
self.indices = zeros(max_tris * 3, dtype=warp.int32, device=self.device)
|
|
5136
|
-
|
|
5137
|
-
# alloc surfacer
|
|
5138
|
-
self.id = ctypes.c_uint64(self.alloc(self.device.context))
|
|
5139
|
-
|
|
5140
|
-
def __del__(self):
|
|
5141
|
-
if not self.id:
|
|
5142
|
-
return
|
|
5143
|
-
|
|
5144
|
-
# use CUDA context guard to avoid side effects during garbage collection
|
|
5145
|
-
with self.device.context_guard:
|
|
5146
|
-
# destroy surfacer
|
|
5147
|
-
self.free(self.id)
|
|
5148
|
-
|
|
5149
|
-
def resize(self, nx: int, ny: int, nz: int, max_verts: int, max_tris: int) -> None:
|
|
5150
|
-
"""Update the expected input and maximum output sizes for the marching cubes calculation.
|
|
5151
|
-
|
|
5152
|
-
This function has no immediate effect on the underlying buffers.
|
|
5153
|
-
The new values take effect on the next :func:`surface` call.
|
|
5154
|
-
|
|
5155
|
-
Args:
|
|
5156
|
-
nx: Number of cubes in the x-direction.
|
|
5157
|
-
ny: Number of cubes in the y-direction.
|
|
5158
|
-
nz: Number of cubes in the z-direction.
|
|
5159
|
-
max_verts: Maximum expected number of vertices (used for array preallocation).
|
|
5160
|
-
max_tris: Maximum expected number of triangles (used for array preallocation).
|
|
5161
|
-
"""
|
|
5162
|
-
# actual allocations will be resized on next call to surface()
|
|
5163
|
-
self.nx = nx
|
|
5164
|
-
self.ny = ny
|
|
5165
|
-
self.nz = nz
|
|
5166
|
-
self.max_verts = max_verts
|
|
5167
|
-
self.max_tris = max_tris
|
|
5168
|
-
|
|
5169
|
-
def surface(self, field: array(dtype=float, ndim=3), threshold: float) -> None:
|
|
5170
|
-
"""Compute a 2D surface mesh of a given isosurface from a 3D scalar field.
|
|
5171
|
-
|
|
5172
|
-
The triangles and vertices defining the output mesh are written to the
|
|
5173
|
-
:attr:`indices` and :attr:`verts` arrays.
|
|
5174
|
-
|
|
5175
|
-
Args:
|
|
5176
|
-
field: Scalar field from which to generate a mesh.
|
|
5177
|
-
threshold: Target isosurface value.
|
|
5178
|
-
|
|
5179
|
-
Raises:
|
|
5180
|
-
ValueError: ``field`` is not a 3D array.
|
|
5181
|
-
ValueError: Marching cubes shape does not match the shape of ``field``.
|
|
5182
|
-
RuntimeError: :attr:`max_verts` and/or :attr:`max_tris` might be too small to hold the surface mesh.
|
|
5183
|
-
"""
|
|
5184
|
-
|
|
5185
|
-
# WP_API int marching_cubes_surface_host(const float* field, int nx, int ny, int nz, float threshold, wp::vec3* verts, int* triangles, int max_verts, int max_tris, int* out_num_verts, int* out_num_tris);
|
|
5186
|
-
num_verts = ctypes.c_int(0)
|
|
5187
|
-
num_tris = ctypes.c_int(0)
|
|
5188
|
-
|
|
5189
|
-
self.runtime.core.marching_cubes_surface_device.restype = ctypes.c_int
|
|
5190
|
-
|
|
5191
|
-
# For now we require that input field shape matches nx, ny, nz
|
|
5192
|
-
if field.ndim != 3:
|
|
5193
|
-
raise ValueError(f"Input field must be a three-dimensional array (got {field.ndim}).")
|
|
5194
|
-
if field.shape[0] != self.nx or field.shape[1] != self.ny or field.shape[2] != self.nz:
|
|
5195
|
-
raise ValueError(
|
|
5196
|
-
f"Marching cubes shape ({self.nx}, {self.ny}, {self.nz}) does not match the "
|
|
5197
|
-
f"input array shape {field.shape}."
|
|
5198
|
-
)
|
|
5199
|
-
|
|
5200
|
-
error = self.runtime.core.marching_cubes_surface_device(
|
|
5201
|
-
self.id,
|
|
5202
|
-
ctypes.cast(field.ptr, ctypes.c_void_p),
|
|
5203
|
-
self.nx,
|
|
5204
|
-
self.ny,
|
|
5205
|
-
self.nz,
|
|
5206
|
-
ctypes.c_float(threshold),
|
|
5207
|
-
ctypes.cast(self.verts.ptr, ctypes.c_void_p),
|
|
5208
|
-
ctypes.cast(self.indices.ptr, ctypes.c_void_p),
|
|
5209
|
-
self.max_verts,
|
|
5210
|
-
self.max_tris,
|
|
5211
|
-
ctypes.c_void_p(ctypes.addressof(num_verts)),
|
|
5212
|
-
ctypes.c_void_p(ctypes.addressof(num_tris)),
|
|
5213
|
-
)
|
|
5214
|
-
|
|
5215
|
-
if error:
|
|
5216
|
-
raise RuntimeError(
|
|
5217
|
-
f"Buffers may not be large enough, marching cubes required at least {num_verts} vertices, and {num_tris} triangles."
|
|
5218
|
-
)
|
|
5219
|
-
|
|
5220
|
-
# resize the geometry arrays
|
|
5221
|
-
self.verts.shape = (num_verts.value,)
|
|
5222
|
-
self.indices.shape = (num_tris.value * 3,)
|
|
5223
|
-
|
|
5224
|
-
self.verts.size = num_verts.value
|
|
5225
|
-
self.indices.size = num_tris.value * 3
|
|
5515
|
+
self.runtime.core.wp_hash_grid_destroy_device(self.id)
|
|
5226
5516
|
|
|
5227
5517
|
|
|
5228
5518
|
generic_types = (Any, Scalar, Float, Int)
|