warp-lang 1.8.0__py3-none-win_amd64.whl → 1.9.0__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.

Files changed (153) hide show
  1. warp/__init__.py +282 -103
  2. warp/__init__.pyi +482 -110
  3. warp/bin/warp-clang.dll +0 -0
  4. warp/bin/warp.dll +0 -0
  5. warp/build.py +93 -30
  6. warp/build_dll.py +48 -63
  7. warp/builtins.py +955 -137
  8. warp/codegen.py +327 -209
  9. warp/config.py +1 -1
  10. warp/context.py +1363 -800
  11. warp/examples/core/example_marching_cubes.py +1 -0
  12. warp/examples/core/example_render_opengl.py +100 -3
  13. warp/examples/fem/example_apic_fluid.py +98 -52
  14. warp/examples/fem/example_convection_diffusion_dg.py +25 -4
  15. warp/examples/fem/example_diffusion_mgpu.py +8 -3
  16. warp/examples/fem/utils.py +68 -22
  17. warp/examples/interop/example_jax_callable.py +34 -4
  18. warp/examples/interop/example_jax_kernel.py +27 -1
  19. warp/fabric.py +1 -1
  20. warp/fem/cache.py +27 -19
  21. warp/fem/domain.py +2 -2
  22. warp/fem/field/nodal_field.py +2 -2
  23. warp/fem/field/virtual.py +266 -166
  24. warp/fem/geometry/geometry.py +5 -5
  25. warp/fem/integrate.py +200 -91
  26. warp/fem/space/restriction.py +4 -0
  27. warp/fem/space/shape/tet_shape_function.py +3 -10
  28. warp/jax_experimental/custom_call.py +1 -1
  29. warp/jax_experimental/ffi.py +203 -54
  30. warp/marching_cubes.py +708 -0
  31. warp/native/array.h +103 -8
  32. warp/native/builtin.h +90 -9
  33. warp/native/bvh.cpp +64 -28
  34. warp/native/bvh.cu +58 -58
  35. warp/native/bvh.h +2 -2
  36. warp/native/clang/clang.cpp +7 -7
  37. warp/native/coloring.cpp +13 -3
  38. warp/native/crt.cpp +2 -2
  39. warp/native/crt.h +3 -5
  40. warp/native/cuda_util.cpp +42 -11
  41. warp/native/cuda_util.h +10 -4
  42. warp/native/exports.h +1842 -1908
  43. warp/native/fabric.h +2 -1
  44. warp/native/hashgrid.cpp +37 -37
  45. warp/native/hashgrid.cu +2 -2
  46. warp/native/initializer_array.h +1 -1
  47. warp/native/intersect.h +4 -4
  48. warp/native/mat.h +1913 -119
  49. warp/native/mathdx.cpp +43 -43
  50. warp/native/mesh.cpp +24 -24
  51. warp/native/mesh.cu +26 -26
  52. warp/native/mesh.h +5 -3
  53. warp/native/nanovdb/GridHandle.h +179 -12
  54. warp/native/nanovdb/HostBuffer.h +8 -7
  55. warp/native/nanovdb/NanoVDB.h +517 -895
  56. warp/native/nanovdb/NodeManager.h +323 -0
  57. warp/native/nanovdb/PNanoVDB.h +2 -2
  58. warp/native/quat.h +337 -16
  59. warp/native/rand.h +7 -7
  60. warp/native/range.h +7 -1
  61. warp/native/reduce.cpp +10 -10
  62. warp/native/reduce.cu +13 -14
  63. warp/native/runlength_encode.cpp +2 -2
  64. warp/native/runlength_encode.cu +5 -5
  65. warp/native/scan.cpp +3 -3
  66. warp/native/scan.cu +4 -4
  67. warp/native/sort.cpp +10 -10
  68. warp/native/sort.cu +22 -22
  69. warp/native/sparse.cpp +8 -8
  70. warp/native/sparse.cu +14 -14
  71. warp/native/spatial.h +366 -17
  72. warp/native/svd.h +23 -8
  73. warp/native/temp_buffer.h +2 -2
  74. warp/native/tile.h +303 -70
  75. warp/native/tile_radix_sort.h +5 -1
  76. warp/native/tile_reduce.h +16 -25
  77. warp/native/tuple.h +2 -2
  78. warp/native/vec.h +385 -18
  79. warp/native/volume.cpp +54 -54
  80. warp/native/volume.cu +1 -1
  81. warp/native/volume.h +2 -1
  82. warp/native/volume_builder.cu +30 -37
  83. warp/native/warp.cpp +150 -149
  84. warp/native/warp.cu +337 -193
  85. warp/native/warp.h +227 -226
  86. warp/optim/linear.py +736 -271
  87. warp/render/imgui_manager.py +289 -0
  88. warp/render/render_opengl.py +137 -57
  89. warp/render/render_usd.py +0 -1
  90. warp/sim/collide.py +1 -2
  91. warp/sim/graph_coloring.py +2 -2
  92. warp/sim/integrator_vbd.py +10 -2
  93. warp/sparse.py +559 -176
  94. warp/tape.py +2 -0
  95. warp/tests/aux_test_module_aot.py +7 -0
  96. warp/tests/cuda/test_async.py +3 -3
  97. warp/tests/cuda/test_conditional_captures.py +101 -0
  98. warp/tests/geometry/test_marching_cubes.py +233 -12
  99. warp/tests/sim/test_cloth.py +89 -6
  100. warp/tests/sim/test_coloring.py +82 -7
  101. warp/tests/test_array.py +56 -5
  102. warp/tests/test_assert.py +53 -0
  103. warp/tests/test_atomic_cas.py +127 -114
  104. warp/tests/test_codegen.py +3 -2
  105. warp/tests/test_context.py +8 -15
  106. warp/tests/test_enum.py +136 -0
  107. warp/tests/test_examples.py +2 -2
  108. warp/tests/test_fem.py +45 -2
  109. warp/tests/test_fixedarray.py +229 -0
  110. warp/tests/test_func.py +18 -15
  111. warp/tests/test_future_annotations.py +7 -5
  112. warp/tests/test_linear_solvers.py +30 -0
  113. warp/tests/test_map.py +1 -1
  114. warp/tests/test_mat.py +1540 -378
  115. warp/tests/test_mat_assign_copy.py +178 -0
  116. warp/tests/test_mat_constructors.py +574 -0
  117. warp/tests/test_module_aot.py +287 -0
  118. warp/tests/test_print.py +69 -0
  119. warp/tests/test_quat.py +162 -34
  120. warp/tests/test_quat_assign_copy.py +145 -0
  121. warp/tests/test_reload.py +2 -1
  122. warp/tests/test_sparse.py +103 -0
  123. warp/tests/test_spatial.py +140 -34
  124. warp/tests/test_spatial_assign_copy.py +160 -0
  125. warp/tests/test_static.py +48 -0
  126. warp/tests/test_struct.py +43 -3
  127. warp/tests/test_tape.py +38 -0
  128. warp/tests/test_types.py +0 -20
  129. warp/tests/test_vec.py +216 -441
  130. warp/tests/test_vec_assign_copy.py +143 -0
  131. warp/tests/test_vec_constructors.py +325 -0
  132. warp/tests/tile/test_tile.py +206 -152
  133. warp/tests/tile/test_tile_cholesky.py +605 -0
  134. warp/tests/tile/test_tile_load.py +169 -0
  135. warp/tests/tile/test_tile_mathdx.py +2 -558
  136. warp/tests/tile/test_tile_matmul.py +179 -0
  137. warp/tests/tile/test_tile_mlp.py +1 -1
  138. warp/tests/tile/test_tile_reduce.py +100 -11
  139. warp/tests/tile/test_tile_shared_memory.py +16 -16
  140. warp/tests/tile/test_tile_sort.py +59 -55
  141. warp/tests/unittest_suites.py +16 -0
  142. warp/tests/walkthrough_debug.py +1 -1
  143. warp/thirdparty/unittest_parallel.py +108 -9
  144. warp/types.py +554 -264
  145. warp/utils.py +68 -86
  146. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/METADATA +28 -65
  147. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/RECORD +150 -138
  148. warp/native/marching.cpp +0 -19
  149. warp/native/marching.cu +0 -514
  150. warp/native/marching.h +0 -19
  151. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/WHEEL +0 -0
  152. {warp_lang-1.8.0.dist-info → warp_lang-1.9.0.dist-info}/licenses/LICENSE.md +0 -0
  153. {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.float_to_half_bits(value)
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.half_bits_to_float(value)
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
- return [vec_t.scalar_export(x) for x in super().__getitem__(key)]
333
+ values = tuple(vec_t.scalar_export(x) for x in super().__getitem__(key))
264
334
  else:
265
- return super().__getitem__(key)
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(self, y)
406
+ return _binary_op(self, warp.add, y, vec_t)
329
407
 
330
408
  def __radd__(self, y):
331
- return warp.add(y, self)
409
+ return _rbinary_op(self, warp.add, y, vec_t)
332
410
 
333
411
  def __sub__(self, y):
334
- return warp.sub(self, y)
412
+ return _binary_op(self, warp.sub, y, vec_t)
335
413
 
336
414
  def __rsub__(self, y):
337
- return warp.sub(y, self)
415
+ return _rbinary_op(self, warp.sub, y, vec_t)
338
416
 
339
417
  def __mul__(self, y):
340
- return warp.mul(self, y)
418
+ return _binary_op(self, warp.mul, y, vec_t, cw=False)
341
419
 
342
420
  def __rmul__(self, x):
343
- return warp.mul(x, self)
421
+ return _rbinary_op(self, warp.mul, x, vec_t, cw=False)
344
422
 
345
423
  def __truediv__(self, y):
346
- return warp.div(self, y)
424
+ return _binary_op(self, warp.div, y, vec_t, cw=False)
347
425
 
348
426
  def __rtruediv__(self, x):
349
- return warp.div(x, self)
427
+ return _rbinary_op(self, warp.div, x, vec_t, cw=False)
350
428
 
351
429
  def __mod__(self, x):
352
- return warp.mod(self, x)
430
+ return _binary_op(self, warp.mod, x, vec_t)
353
431
 
354
432
  def __rmod__(self, x):
355
- return warp.mod(x, self)
433
+ return _rbinary_op(self, warp.mod, x, vec_t)
356
434
 
357
435
  def __pos__(self):
358
- return warp.pos(self)
436
+ return _unary_op(self, warp.pos, vec_t)
359
437
 
360
438
  def __neg__(self):
361
- return warp.neg(self)
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(self, y)
549
+ return _binary_op(self, warp.add, y, mat_t)
471
550
 
472
551
  def __radd__(self, y):
473
- return warp.add(y, self)
552
+ return _rbinary_op(self, warp.add, y, mat_t)
474
553
 
475
554
  def __sub__(self, y):
476
- return warp.sub(self, y)
555
+ return _binary_op(self, warp.sub, y, mat_t)
477
556
 
478
557
  def __rsub__(self, y):
479
- return warp.sub(y, self)
558
+ return _rbinary_op(self, warp.sub, y, mat_t)
480
559
 
481
560
  def __mul__(self, y):
482
- return warp.mul(self, y)
561
+ return _binary_op(self, warp.mul, y, mat_t, cw=False)
483
562
 
484
563
  def __rmul__(self, x):
485
- return warp.mul(x, self)
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(self, y)
573
+ return _binary_op(self, warp.div, y, mat_t, cw=False)
495
574
 
496
575
  def __rtruediv__(self, x):
497
- return warp.div(x, self)
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(self)
585
+ return _unary_op(self, warp.pos, mat_t)
501
586
 
502
587
  def __neg__(self):
503
- return warp.neg(self)
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
- if any(isinstance(x, slice) for x in key):
567
- raise KeyError("Slices are not supported when indexing matrices using the `m[i, j]` notation")
568
- return mat_t.scalar_export(super().__getitem__(key[0] * self._shape_[1] + key[1]))
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
- if any(isinstance(x, slice) for x in key):
581
- raise KeyError("Slices are not supported when indexing matrices using the `m[i, j]` notation")
582
- try:
583
- return super().__setitem__(key[0] * self._shape_[1] + key[1], mat_t.scalar_import(value))
584
- except (TypeError, ctypes.ArgumentError):
585
- raise TypeError(
586
- f"Expected to assign a `{self._wp_scalar_type_.__name__}` value "
587
- f"but got `{type(value).__name__}` instead"
588
- ) from None
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
- raise KeyError("Slices are not supported when indexing matrices using the `m[start:end]` notation")
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
- pass
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"array(ndim={t.ndim}, dtype={type_repr(t.dtype)})"
1531
- return f"array(shape={t.shape}, dtype={type_repr(t.dtype)})"
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 is_array(a) and type(a) is type(b) and types_equal(a.dtype, b.dtype, match_generic=match_generic):
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.cuda_stream_wait_stream(
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
- # Note: we use a shared dict for all array instances
2726
- if array._vars is None:
2727
- array._vars = {"shape": warp.codegen.Var("shape", shape_t)}
2728
- return array._vars
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.array_fill_device(
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.array_fill_host(carr_ptr, ARRAY_TYPE_REGULAR, cvalue_ptr, cvalue_size)
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.cuda_ipc_get_mem_handle(self.ptr, ipc_handle_buffer)
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.cuda_ipc_close_mem_handle(ptr)
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.cuda_ipc_open_mem_handle(device.context, handle)
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.array_fill_device(
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.array_fill_host(ctype_ptr, self.type_id, cvalue_ptr, cvalue_size)
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.bvh_create_host(
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.bvh_create_device(
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.bvh_destroy_host(self.id)
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.bvh_destroy_device(self.id)
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.bvh_refit_host(self.id)
4137
+ self.runtime.core.wp_bvh_refit_host(self.id)
3697
4138
  else:
3698
- self.runtime.core.bvh_refit_device(self.id)
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.mesh_create_host(
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.mesh_create_device(
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.mesh_destroy_host(self.id)
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.mesh_destroy_device(self.id)
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.mesh_refit_host(self.id)
4260
+ self.runtime.core.wp_mesh_refit_host(self.id)
3820
4261
  else:
3821
- self.runtime.core.mesh_refit_device(self.id)
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.mesh_set_points_host(self.id, points_new.__ctype__())
4292
+ self.runtime.core.wp_mesh_set_points_host(self.id, points_new.__ctype__())
3852
4293
  else:
3853
- self.runtime.core.mesh_set_points_device(self.id, points_new.__ctype__())
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.mesh_set_velocities_host(self.id, velocities_new.__ctype__())
4322
+ self.runtime.core.wp_mesh_set_velocities_host(self.id, velocities_new.__ctype__())
3882
4323
  else:
3883
- self.runtime.core.mesh_set_velocities_device(self.id, velocities_new.__ctype__())
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.volume_create_host(
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.volume_create_device(
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.volume_destroy_host(self.id)
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.volume_destroy_device(self.id)
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.volume_get_buffer_info(self.id, ctypes.byref(buf), ctypes.byref(size))
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.volume_get_tile_and_voxel_count(self.id, ctypes.byref(tile_count), ctypes.byref(voxel_count))
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.volume_get_tiles_host(self.id, out.ptr)
4422
+ self.runtime.core.wp_volume_get_tiles_host(self.id, out.ptr)
3980
4423
  else:
3981
- self.runtime.core.volume_get_tiles_device(self.id, out.ptr)
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.volume_get_tile_and_voxel_count(self.id, ctypes.byref(tile_count), ctypes.byref(voxel_count))
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.volume_get_voxels_host(self.id, out.ptr)
4463
+ self.runtime.core.wp_volume_get_voxels_host(self.id, out.ptr)
4019
4464
  else:
4020
- self.runtime.core.volume_get_voxels_device(self.id, out.ptr)
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.volume_get_voxel_size(self.id, ctypes.byref(dx), ctypes.byref(dy), ctypes.byref(dz))
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.volume_get_grid_info(
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.volume_get_blind_data_count(self.id)
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.volume_get_blind_data_info(
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.volume_get_buffer_info(
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.volume_index_from_tiles_device(
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.volume_from_tiles_device(
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.volume_from_active_voxels_device(
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.hash_grid_create_host(dim_x, dim_y, dim_z)
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.hash_grid_create_device(self.device.context, dim_x, dim_y, dim_z)
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.hash_grid_update_host(self.id, radius, ctypes.byref(points.__ctype__()))
5494
+ self.runtime.core.wp_hash_grid_update_host(self.id, radius, ctypes.byref(points.__ctype__()))
5050
5495
  else:
5051
- self.runtime.core.hash_grid_update_device(self.id, radius, ctypes.byref(points.__ctype__()))
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.hash_grid_reserve_host(self.id, num_points)
5501
+ self.runtime.core.wp_hash_grid_reserve_host(self.id, num_points)
5057
5502
  else:
5058
- self.runtime.core.hash_grid_reserve_device(self.id, num_points)
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.hash_grid_destroy_host(self.id)
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.hash_grid_destroy_device(self.id)
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)