warp-lang 1.6.2__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.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 (179) hide show
  1. warp/__init__.py +7 -1
  2. warp/bin/libwarp-clang.dylib +0 -0
  3. warp/bin/libwarp.dylib +0 -0
  4. warp/build.py +410 -0
  5. warp/build_dll.py +6 -14
  6. warp/builtins.py +452 -362
  7. warp/codegen.py +179 -119
  8. warp/config.py +42 -6
  9. warp/context.py +490 -271
  10. warp/dlpack.py +8 -6
  11. warp/examples/assets/nonuniform.usd +0 -0
  12. warp/examples/assets/nvidia_logo.png +0 -0
  13. warp/examples/benchmarks/benchmark_tile_load_store.py +103 -0
  14. warp/examples/core/example_sample_mesh.py +300 -0
  15. warp/examples/fem/example_apic_fluid.py +1 -1
  16. warp/examples/fem/example_burgers.py +2 -2
  17. warp/examples/fem/example_deformed_geometry.py +1 -1
  18. warp/examples/fem/example_distortion_energy.py +1 -1
  19. warp/examples/fem/example_magnetostatics.py +6 -6
  20. warp/examples/fem/utils.py +9 -3
  21. warp/examples/interop/example_jax_callable.py +116 -0
  22. warp/examples/interop/example_jax_ffi_callback.py +132 -0
  23. warp/examples/interop/example_jax_kernel.py +205 -0
  24. warp/examples/optim/example_fluid_checkpoint.py +497 -0
  25. warp/examples/tile/example_tile_matmul.py +2 -4
  26. warp/fem/__init__.py +11 -1
  27. warp/fem/adaptivity.py +4 -4
  28. warp/fem/field/nodal_field.py +22 -68
  29. warp/fem/field/virtual.py +62 -23
  30. warp/fem/geometry/adaptive_nanogrid.py +9 -10
  31. warp/fem/geometry/closest_point.py +1 -1
  32. warp/fem/geometry/deformed_geometry.py +5 -2
  33. warp/fem/geometry/geometry.py +5 -0
  34. warp/fem/geometry/grid_2d.py +12 -12
  35. warp/fem/geometry/grid_3d.py +12 -15
  36. warp/fem/geometry/hexmesh.py +5 -7
  37. warp/fem/geometry/nanogrid.py +9 -11
  38. warp/fem/geometry/quadmesh.py +13 -13
  39. warp/fem/geometry/tetmesh.py +3 -4
  40. warp/fem/geometry/trimesh.py +3 -8
  41. warp/fem/integrate.py +262 -93
  42. warp/fem/linalg.py +5 -5
  43. warp/fem/quadrature/pic_quadrature.py +37 -22
  44. warp/fem/quadrature/quadrature.py +194 -25
  45. warp/fem/space/__init__.py +1 -1
  46. warp/fem/space/basis_function_space.py +4 -2
  47. warp/fem/space/basis_space.py +25 -18
  48. warp/fem/space/hexmesh_function_space.py +2 -2
  49. warp/fem/space/partition.py +6 -2
  50. warp/fem/space/quadmesh_function_space.py +8 -8
  51. warp/fem/space/shape/cube_shape_function.py +23 -23
  52. warp/fem/space/shape/square_shape_function.py +12 -12
  53. warp/fem/space/shape/triangle_shape_function.py +1 -1
  54. warp/fem/space/tetmesh_function_space.py +3 -3
  55. warp/fem/space/trimesh_function_space.py +2 -2
  56. warp/fem/utils.py +12 -6
  57. warp/jax.py +14 -1
  58. warp/jax_experimental/__init__.py +16 -0
  59. warp/{jax_experimental.py → jax_experimental/custom_call.py} +14 -27
  60. warp/jax_experimental/ffi.py +698 -0
  61. warp/jax_experimental/xla_ffi.py +602 -0
  62. warp/math.py +89 -0
  63. warp/native/array.h +13 -0
  64. warp/native/builtin.h +29 -3
  65. warp/native/bvh.cpp +3 -1
  66. warp/native/bvh.cu +42 -14
  67. warp/native/bvh.h +2 -1
  68. warp/native/clang/clang.cpp +30 -3
  69. warp/native/cuda_util.cpp +14 -0
  70. warp/native/cuda_util.h +2 -0
  71. warp/native/exports.h +68 -63
  72. warp/native/intersect.h +26 -26
  73. warp/native/intersect_adj.h +33 -33
  74. warp/native/marching.cu +1 -1
  75. warp/native/mat.h +513 -9
  76. warp/native/mesh.h +10 -10
  77. warp/native/quat.h +99 -11
  78. warp/native/rand.h +6 -0
  79. warp/native/sort.cpp +122 -59
  80. warp/native/sort.cu +152 -15
  81. warp/native/sort.h +8 -1
  82. warp/native/sparse.cpp +43 -22
  83. warp/native/sparse.cu +52 -17
  84. warp/native/svd.h +116 -0
  85. warp/native/tile.h +301 -105
  86. warp/native/tile_reduce.h +46 -3
  87. warp/native/vec.h +68 -7
  88. warp/native/volume.cpp +85 -113
  89. warp/native/volume_builder.cu +25 -10
  90. warp/native/volume_builder.h +6 -0
  91. warp/native/warp.cpp +5 -6
  92. warp/native/warp.cu +99 -10
  93. warp/native/warp.h +19 -10
  94. warp/optim/linear.py +10 -10
  95. warp/sim/articulation.py +4 -4
  96. warp/sim/collide.py +21 -10
  97. warp/sim/import_mjcf.py +449 -155
  98. warp/sim/import_urdf.py +32 -12
  99. warp/sim/integrator_euler.py +5 -5
  100. warp/sim/integrator_featherstone.py +3 -10
  101. warp/sim/integrator_vbd.py +207 -2
  102. warp/sim/integrator_xpbd.py +5 -5
  103. warp/sim/model.py +42 -13
  104. warp/sim/utils.py +2 -2
  105. warp/sparse.py +642 -555
  106. warp/stubs.py +216 -19
  107. warp/tests/__main__.py +0 -15
  108. warp/tests/cuda/__init__.py +0 -0
  109. warp/tests/{test_mempool.py → cuda/test_mempool.py} +39 -0
  110. warp/tests/{test_streams.py → cuda/test_streams.py} +71 -0
  111. warp/tests/geometry/__init__.py +0 -0
  112. warp/tests/{test_mesh_query_point.py → geometry/test_mesh_query_point.py} +66 -63
  113. warp/tests/{test_mesh_query_ray.py → geometry/test_mesh_query_ray.py} +1 -1
  114. warp/tests/{test_volume.py → geometry/test_volume.py} +41 -6
  115. warp/tests/interop/__init__.py +0 -0
  116. warp/tests/{test_dlpack.py → interop/test_dlpack.py} +28 -5
  117. warp/tests/sim/__init__.py +0 -0
  118. warp/tests/{disabled_kinematics.py → sim/disabled_kinematics.py} +9 -10
  119. warp/tests/{test_collision.py → sim/test_collision.py} +2 -2
  120. warp/tests/{test_model.py → sim/test_model.py} +40 -0
  121. warp/tests/{test_sim_kinematics.py → sim/test_sim_kinematics.py} +2 -1
  122. warp/tests/sim/test_vbd.py +597 -0
  123. warp/tests/test_bool.py +1 -1
  124. warp/tests/test_examples.py +28 -36
  125. warp/tests/test_fem.py +23 -4
  126. warp/tests/test_linear_solvers.py +0 -11
  127. warp/tests/test_mat.py +233 -79
  128. warp/tests/test_mat_scalar_ops.py +4 -4
  129. warp/tests/test_overwrite.py +0 -60
  130. warp/tests/test_quat.py +67 -46
  131. warp/tests/test_rand.py +44 -37
  132. warp/tests/test_sparse.py +47 -6
  133. warp/tests/test_spatial.py +75 -0
  134. warp/tests/test_static.py +1 -1
  135. warp/tests/test_utils.py +84 -4
  136. warp/tests/test_vec.py +46 -34
  137. warp/tests/tile/__init__.py +0 -0
  138. warp/tests/{test_tile.py → tile/test_tile.py} +136 -51
  139. warp/tests/{test_tile_load.py → tile/test_tile_load.py} +1 -1
  140. warp/tests/{test_tile_mathdx.py → tile/test_tile_mathdx.py} +9 -6
  141. warp/tests/{test_tile_mlp.py → tile/test_tile_mlp.py} +25 -14
  142. warp/tests/{test_tile_reduce.py → tile/test_tile_reduce.py} +60 -1
  143. warp/tests/{test_tile_view.py → tile/test_tile_view.py} +1 -1
  144. warp/tests/unittest_serial.py +1 -0
  145. warp/tests/unittest_suites.py +45 -59
  146. warp/tests/unittest_utils.py +2 -1
  147. warp/thirdparty/unittest_parallel.py +3 -1
  148. warp/types.py +110 -658
  149. warp/utils.py +137 -72
  150. {warp_lang-1.6.2.dist-info → warp_lang-1.7.0.dist-info}/METADATA +29 -7
  151. {warp_lang-1.6.2.dist-info → warp_lang-1.7.0.dist-info}/RECORD +172 -162
  152. {warp_lang-1.6.2.dist-info → warp_lang-1.7.0.dist-info}/WHEEL +1 -1
  153. warp/examples/optim/example_walker.py +0 -317
  154. warp/native/cutlass_gemm.cpp +0 -43
  155. warp/native/cutlass_gemm.cu +0 -382
  156. warp/tests/test_matmul.py +0 -511
  157. warp/tests/test_matmul_lite.py +0 -411
  158. warp/tests/test_vbd.py +0 -386
  159. warp/tests/unused_test_misc.py +0 -77
  160. /warp/tests/{test_async.py → cuda/test_async.py} +0 -0
  161. /warp/tests/{test_ipc.py → cuda/test_ipc.py} +0 -0
  162. /warp/tests/{test_multigpu.py → cuda/test_multigpu.py} +0 -0
  163. /warp/tests/{test_peer.py → cuda/test_peer.py} +0 -0
  164. /warp/tests/{test_pinned.py → cuda/test_pinned.py} +0 -0
  165. /warp/tests/{test_bvh.py → geometry/test_bvh.py} +0 -0
  166. /warp/tests/{test_hash_grid.py → geometry/test_hash_grid.py} +0 -0
  167. /warp/tests/{test_marching_cubes.py → geometry/test_marching_cubes.py} +0 -0
  168. /warp/tests/{test_mesh.py → geometry/test_mesh.py} +0 -0
  169. /warp/tests/{test_mesh_query_aabb.py → geometry/test_mesh_query_aabb.py} +0 -0
  170. /warp/tests/{test_volume_write.py → geometry/test_volume_write.py} +0 -0
  171. /warp/tests/{test_jax.py → interop/test_jax.py} +0 -0
  172. /warp/tests/{test_paddle.py → interop/test_paddle.py} +0 -0
  173. /warp/tests/{test_torch.py → interop/test_torch.py} +0 -0
  174. /warp/tests/{flaky_test_sim_grad.py → sim/flaky_test_sim_grad.py} +0 -0
  175. /warp/tests/{test_coloring.py → sim/test_coloring.py} +0 -0
  176. /warp/tests/{test_sim_grad_bounce_linear.py → sim/test_sim_grad_bounce_linear.py} +0 -0
  177. /warp/tests/{test_tile_shared_memory.py → tile/test_tile_shared_memory.py} +0 -0
  178. {warp_lang-1.6.2.dist-info → warp_lang-1.7.0.dist-info/licenses}/LICENSE.md +0 -0
  179. {warp_lang-1.6.2.dist-info → warp_lang-1.7.0.dist-info}/top_level.txt +0 -0
@@ -22,11 +22,6 @@ import warp.examples
22
22
  import warp.optim
23
23
  from warp.tests.unittest_utils import *
24
24
 
25
- wp.init()
26
-
27
- # needs to be constant for the whole module
28
- NUM_THREADS = 32
29
-
30
25
 
31
26
  def create_layer(rng, dim_in, dim_hid, dtype=float):
32
27
  w = rng.uniform(-1.0 / np.sqrt(dim_in), 1.0 / np.sqrt(dim_in), (dim_hid, dim_in))
@@ -45,10 +40,12 @@ def create_array(rng, dim_in, dim_hid, dtype=float):
45
40
  return a
46
41
 
47
42
 
48
- @unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
49
43
  def test_multi_layer_nn(test, device):
50
44
  import torch as tc
51
45
 
46
+ if device.is_cuda and not wp.context.runtime.core.is_mathdx_enabled():
47
+ test.skipTest("Skipping test on CUDA device without MathDx (tolerance)")
48
+
52
49
  NUM_FREQ = wp.constant(8)
53
50
 
54
51
  DIM_IN = wp.constant(4 * NUM_FREQ) # sin,cos for both x,y at each frequency
@@ -60,7 +57,13 @@ def test_multi_layer_nn(test, device):
60
57
 
61
58
  BATCH_SIZE = min(512, int((IMG_WIDTH * IMG_HEIGHT) / 8))
62
59
 
60
+ if device.is_cpu:
61
+ NUM_THREADS = 1
62
+ else:
63
+ NUM_THREADS = 32
64
+
63
65
  dtype = wp.float16
66
+ npdtype = wp.types.warp_type_to_np_dtype[dtype]
64
67
 
65
68
  @wp.func
66
69
  def relu(x: dtype):
@@ -74,7 +77,7 @@ def test_multi_layer_nn(test, device):
74
77
  def zero(loss: wp.array(dtype=float)):
75
78
  loss[0] = 0.0
76
79
 
77
- @wp.kernel
80
+ @wp.kernel(module="unique")
78
81
  def compute(
79
82
  batches: wp.array(dtype=int),
80
83
  input: wp.array2d(dtype=dtype),
@@ -170,7 +173,9 @@ def test_multi_layer_nn(test, device):
170
173
  input = create_array(rng, IMG_WIDTH * IMG_HEIGHT, DIM_IN, dtype=dtype)
171
174
  output = create_array(rng, IMG_WIDTH * IMG_HEIGHT, DIM_OUT)
172
175
 
173
- reference_np = np.load(os.path.join(os.path.dirname(__file__), "assets/pixel.npy"), allow_pickle=True) / 255.0
176
+ reference_np = (
177
+ np.load(os.path.join(os.path.dirname(__file__), "..", "assets", "pixel.npy"), allow_pickle=True) / 255.0
178
+ )
174
179
  reference = wp.array(reference_np, dtype=float)
175
180
 
176
181
  assert reference.shape[1] == IMG_WIDTH * IMG_HEIGHT
@@ -232,7 +237,7 @@ def test_multi_layer_nn(test, device):
232
237
  z_np = np.maximum(weights_3.numpy() @ z_np + bias_3.numpy(), 0.0)
233
238
 
234
239
  # test numpy forward
235
- assert_np_equal(output.numpy()[:, indices], z_np, tol=1.0e-2)
240
+ assert_np_equal(output.numpy()[:, indices].astype(npdtype), z_np, tol=1.0e-2)
236
241
 
237
242
  # torch
238
243
  input_tc = tc.tensor(input.numpy()[:, indices], requires_grad=True, device=torch_device)
@@ -260,7 +265,9 @@ def test_multi_layer_nn(test, device):
260
265
  l_tc.backward()
261
266
 
262
267
  # test torch
263
- assert_np_equal(z_tc.cpu().detach().numpy(), output.numpy()[:, indices], tol=1.0e-2)
268
+ assert_np_equal(
269
+ z_tc.cpu().detach().numpy(), output.numpy()[:, indices].astype(npdtype), tol=1.0e-2
270
+ )
264
271
  assert_np_equal(weights_0.grad.numpy(), weights_0_tc.grad.cpu().detach().numpy(), tol=1.0e-2)
265
272
  assert_np_equal(bias_0.grad.numpy(), bias_0_tc.grad.cpu().detach().numpy(), tol=1.0e-2)
266
273
  assert_np_equal(weights_1.grad.numpy(), weights_1_tc.grad.cpu().detach().numpy(), tol=1.0e-2)
@@ -277,7 +284,6 @@ def test_multi_layer_nn(test, device):
277
284
  test.assertLess(loss.numpy()[0], 0.002)
278
285
 
279
286
 
280
- @unittest.skipUnless(wp.context.runtime.core.is_mathdx_enabled(), "Warp was not built with MathDx support")
281
287
  def test_single_layer_nn(test, device):
282
288
  import torch as tc
283
289
 
@@ -287,11 +293,16 @@ def test_single_layer_nn(test, device):
287
293
 
288
294
  NUM_BLOCKS = 56
289
295
 
296
+ if device.is_cpu:
297
+ NUM_THREADS = 1
298
+ else:
299
+ NUM_THREADS = 32
300
+
290
301
  @wp.func
291
302
  def relu(x: float):
292
303
  return wp.max(x, 0.0)
293
304
 
294
- @wp.kernel
305
+ @wp.kernel(module="unique")
295
306
  def compute(
296
307
  input: wp.array2d(dtype=float),
297
308
  weights: wp.array2d(dtype=float),
@@ -353,7 +364,6 @@ try:
353
364
  import torch
354
365
 
355
366
  # check which Warp devices work with Torch
356
- # CUDA devices may fail if Torch was not compiled with CUDA support
357
367
  torch_compatible_devices = []
358
368
  torch_compatible_cuda_devices = []
359
369
 
@@ -372,7 +382,7 @@ try:
372
382
  "test_single_layer_nn",
373
383
  test_single_layer_nn,
374
384
  check_output=False,
375
- devices=torch_compatible_cuda_devices,
385
+ devices=torch_compatible_devices,
376
386
  )
377
387
  add_function_test(
378
388
  TestTileMLP,
@@ -388,4 +398,5 @@ except Exception as e:
388
398
 
389
399
  if __name__ == "__main__":
390
400
  wp.clear_kernel_cache()
401
+ wp.clear_lto_cache()
391
402
  unittest.main(verbosity=2, failfast=True)
@@ -176,6 +176,64 @@ def test_tile_reduce_custom(test, device):
176
176
  test.assertAlmostEqual(prod_wp[i], prod_np, places=4)
177
177
 
178
178
 
179
+ @wp.struct
180
+ class KeyValue:
181
+ key: wp.int32
182
+ value: wp.float32
183
+
184
+
185
+ @wp.func
186
+ def kv_max(a: KeyValue, b: KeyValue) -> KeyValue:
187
+ return wp.where(a.value < b.value, b, a)
188
+
189
+
190
+ @wp.kernel
191
+ def initialize_key_value(values: wp.array2d(dtype=wp.float32), keyvalues: wp.array2d(dtype=KeyValue)):
192
+ batch, idx = wp.tid()
193
+ keyvalues[batch, idx] = KeyValue(idx, values[batch, idx])
194
+
195
+
196
+ @wp.kernel(enable_backward=False)
197
+ def tile_reduce_custom_struct_kernel(values: wp.array2d(dtype=KeyValue), res: wp.array(dtype=KeyValue)):
198
+ # output tile index
199
+ i = wp.tid()
200
+
201
+ t = wp.tile_load(values, shape=(1, TILE_DIM), offset=(i, 0))
202
+
203
+ max_el = wp.tile_reduce(kv_max, t)
204
+ wp.tile_store(res, max_el, offset=i)
205
+
206
+
207
+ def test_tile_reduce_custom_struct(test, device):
208
+ batch_count = 56
209
+
210
+ N = TILE_DIM
211
+
212
+ rng = np.random.default_rng(42)
213
+ input = rng.random((batch_count, N), dtype=np.float32)
214
+
215
+ input_wp = wp.array(input, dtype=wp.float32, device=device)
216
+ keyvalues_wp = wp.empty(input_wp.shape, dtype=KeyValue, device=device)
217
+
218
+ wp.launch(initialize_key_value, dim=[batch_count, N], inputs=[input_wp], outputs=[keyvalues_wp], device=device)
219
+
220
+ output_wp = wp.empty(batch_count, dtype=KeyValue, device=device)
221
+
222
+ wp.launch_tiled(
223
+ tile_reduce_custom_struct_kernel,
224
+ dim=[batch_count],
225
+ inputs=[keyvalues_wp],
226
+ outputs=[output_wp],
227
+ block_dim=TILE_DIM,
228
+ device=device,
229
+ )
230
+
231
+ prod_wp = np.array([k for k, v in output_wp.numpy()])
232
+ expected = np.argmax(input, axis=1)
233
+
234
+ assert_np_equal(prod_wp, expected)
235
+
236
+
179
237
  @wp.kernel
180
238
  def tile_grouped_sum_kernel(input: wp.array3d(dtype=float), output: wp.array(dtype=float)):
181
239
  # output tile index
@@ -365,7 +423,7 @@ def test_tile_arange(test, device):
365
423
  assert_np_equal(output.numpy()[4], np.arange(17, 0, -1))
366
424
 
367
425
 
368
- devices = get_cuda_test_devices()
426
+ devices = get_test_devices()
369
427
 
370
428
 
371
429
  class TestTileReduce(unittest.TestCase):
@@ -376,6 +434,7 @@ add_function_test(TestTileReduce, "test_tile_reduce_sum", test_tile_reduce_sum,
376
434
  add_function_test(TestTileReduce, "test_tile_reduce_min", test_tile_reduce_min, devices=devices)
377
435
  add_function_test(TestTileReduce, "test_tile_reduce_max", test_tile_reduce_max, devices=devices)
378
436
  add_function_test(TestTileReduce, "test_tile_reduce_custom", test_tile_reduce_custom, devices=devices)
437
+ add_function_test(TestTileReduce, "test_tile_reduce_custom_struct", test_tile_reduce_custom_struct, devices=devices)
379
438
  add_function_test(TestTileReduce, "test_tile_reduce_grouped_sum", test_tile_reduce_sum, devices=devices)
380
439
  add_function_test(TestTileReduce, "test_tile_reduce_simt", test_tile_reduce_simt, devices=devices)
381
440
  add_function_test(TestTileReduce, "test_tile_ones", test_tile_ones, devices=devices)
@@ -155,7 +155,7 @@ def test_tile_view_offset(test, device):
155
155
  assert_np_equal(a.grad.numpy(), np.ones_like(a.numpy()))
156
156
 
157
157
 
158
- devices = get_cuda_test_devices()
158
+ devices = get_test_devices()
159
159
 
160
160
 
161
161
  class TestTileView(unittest.TestCase):
@@ -23,6 +23,7 @@ def run_suite() -> bool:
23
23
  """Run a test suite"""
24
24
 
25
25
  # force rebuild of all kernels
26
+ wp.clear_lto_cache()
26
27
  wp.clear_kernel_cache()
27
28
  print("Cleared Warp kernel cache")
28
29
 
@@ -95,20 +95,41 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
95
95
 
96
96
  Intended to be modified to create additional test suites
97
97
  """
98
+ from warp.tests.cuda.test_async import TestAsync
99
+ from warp.tests.cuda.test_mempool import TestMempool
100
+ from warp.tests.cuda.test_multigpu import TestMultiGPU
101
+ from warp.tests.cuda.test_peer import TestPeer
102
+ from warp.tests.cuda.test_pinned import TestPinned
103
+ from warp.tests.cuda.test_streams import TestStreams
104
+ from warp.tests.geometry.test_bvh import TestBvh
105
+ from warp.tests.geometry.test_hash_grid import TestHashGrid
106
+ from warp.tests.geometry.test_marching_cubes import TestMarchingCubes
107
+ from warp.tests.geometry.test_mesh import TestMesh
108
+ from warp.tests.geometry.test_mesh_query_aabb import TestMeshQueryAABBMethods
109
+ from warp.tests.geometry.test_mesh_query_point import TestMeshQueryPoint
110
+ from warp.tests.geometry.test_mesh_query_ray import TestMeshQueryRay
111
+ from warp.tests.geometry.test_volume import TestVolume
112
+ from warp.tests.geometry.test_volume_write import TestVolumeWrite
113
+ from warp.tests.interop.test_dlpack import TestDLPack
114
+ from warp.tests.interop.test_jax import TestJax
115
+ from warp.tests.interop.test_torch import TestTorch
116
+ from warp.tests.sim.test_collision import TestCollision
117
+ from warp.tests.sim.test_coloring import TestColoring
118
+ from warp.tests.sim.test_model import TestModel
119
+
120
+ # from warp.tests.test_sim_grad import TestSimGradients Disabled, flaky
121
+ from warp.tests.sim.test_sim_kinematics import TestSimKinematics
122
+ from warp.tests.sim.test_vbd import TestVbd
98
123
  from warp.tests.test_adam import TestAdam
99
124
  from warp.tests.test_arithmetic import TestArithmetic
100
125
  from warp.tests.test_array import TestArray
101
126
  from warp.tests.test_array_reduce import TestArrayReduce
102
- from warp.tests.test_async import TestAsync
103
127
  from warp.tests.test_atomic import TestAtomic
104
128
  from warp.tests.test_bool import TestBool
105
129
  from warp.tests.test_builtins_resolution import TestBuiltinsResolution
106
- from warp.tests.test_bvh import TestBvh
107
130
  from warp.tests.test_closest_point_edge_edge import TestClosestPointEdgeEdgeMethods
108
131
  from warp.tests.test_codegen import TestCodeGen
109
132
  from warp.tests.test_codegen_instancing import TestCodeGenInstancing
110
- from warp.tests.test_collision import TestCollision
111
- from warp.tests.test_coloring import TestColoring
112
133
  from warp.tests.test_compile_consts import TestConstants
113
134
  from warp.tests.test_conditional import TestConditional
114
135
  from warp.tests.test_context import TestContext
@@ -116,7 +137,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
116
137
  from warp.tests.test_ctypes import TestCTypes
117
138
  from warp.tests.test_dense import TestDense
118
139
  from warp.tests.test_devices import TestDevices
119
- from warp.tests.test_dlpack import TestDLPack
120
140
  from warp.tests.test_examples import (
121
141
  TestCoreExamples,
122
142
  TestFemDiffusionExamples,
@@ -134,40 +154,26 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
134
154
  from warp.tests.test_grad import TestGrad
135
155
  from warp.tests.test_grad_customs import TestGradCustoms
136
156
  from warp.tests.test_grad_debug import TestGradDebug
137
- from warp.tests.test_hash_grid import TestHashGrid
138
157
  from warp.tests.test_import import TestImport
139
158
  from warp.tests.test_indexedarray import TestIndexedArray
140
159
  from warp.tests.test_intersect import TestIntersect
141
160
  from warp.tests.test_iter import TestIter
142
- from warp.tests.test_jax import TestJax
143
161
  from warp.tests.test_large import TestLarge
144
162
  from warp.tests.test_launch import TestLaunch
145
163
  from warp.tests.test_lerp import TestLerp
146
164
  from warp.tests.test_linear_solvers import TestLinearSolvers
147
165
  from warp.tests.test_lvalue import TestLValue
148
- from warp.tests.test_marching_cubes import TestMarchingCubes
149
166
  from warp.tests.test_mat import TestMat
150
167
  from warp.tests.test_mat_lite import TestMatLite
151
168
  from warp.tests.test_mat_scalar_ops import TestMatScalarOps
152
169
  from warp.tests.test_math import TestMath
153
- from warp.tests.test_matmul import TestMatmul
154
- from warp.tests.test_matmul_lite import TestMatmulLite
155
- from warp.tests.test_mempool import TestMempool
156
- from warp.tests.test_mesh import TestMesh
157
- from warp.tests.test_mesh_query_aabb import TestMeshQueryAABBMethods
158
- from warp.tests.test_mesh_query_point import TestMeshQueryPoint
159
- from warp.tests.test_mesh_query_ray import TestMeshQueryRay
160
170
  from warp.tests.test_mlp import TestMLP
161
- from warp.tests.test_model import TestModel
162
171
  from warp.tests.test_module_hashing import TestModuleHashing
163
172
  from warp.tests.test_modules_lite import TestModuleLite
164
- from warp.tests.test_multigpu import TestMultiGPU
165
173
  from warp.tests.test_noise import TestNoise
166
174
  from warp.tests.test_operators import TestOperators
167
175
  from warp.tests.test_options import TestOptions
168
176
  from warp.tests.test_overwrite import TestOverwrite
169
- from warp.tests.test_peer import TestPeer
170
- from warp.tests.test_pinned import TestPinned
171
177
  from warp.tests.test_print import TestPrint
172
178
  from warp.tests.test_quat import TestQuat
173
179
  from warp.tests.test_rand import TestRand
@@ -175,34 +181,26 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
175
181
  from warp.tests.test_rounding import TestRounding
176
182
  from warp.tests.test_runlength_encode import TestRunlengthEncode
177
183
  from warp.tests.test_scalar_ops import TestScalarOps
178
-
179
- # from warp.tests.test_sim_grad import TestSimGradients Disabled, flaky
180
- from warp.tests.test_sim_kinematics import TestSimKinematics
181
184
  from warp.tests.test_smoothstep import TestSmoothstep
182
185
  from warp.tests.test_snippet import TestSnippets
183
186
  from warp.tests.test_sparse import TestSparse
184
187
  from warp.tests.test_spatial import TestSpatial
185
188
  from warp.tests.test_special_values import TestSpecialValues
186
189
  from warp.tests.test_static import TestStatic
187
- from warp.tests.test_streams import TestStreams
188
190
  from warp.tests.test_struct import TestStruct
189
191
  from warp.tests.test_tape import TestTape
190
- from warp.tests.test_tile import TestTile
191
- from warp.tests.test_tile_mathdx import TestTileMathDx
192
- from warp.tests.test_tile_reduce import TestTileReduce
193
- from warp.tests.test_tile_shared_memory import TestTileSharedMemory
194
- from warp.tests.test_torch import TestTorch
195
192
  from warp.tests.test_transient_module import TestTransientModule
196
193
  from warp.tests.test_triangle_closest_point import TestTriangleClosestPoint
197
194
  from warp.tests.test_types import TestTypes
198
195
  from warp.tests.test_utils import TestUtils
199
- from warp.tests.test_vbd import TestVBD
200
196
  from warp.tests.test_vec import TestVec
201
197
  from warp.tests.test_vec_lite import TestVecLite
202
198
  from warp.tests.test_vec_scalar_ops import TestVecScalarOps
203
199
  from warp.tests.test_verify_fp import TestVerifyFP
204
- from warp.tests.test_volume import TestVolume
205
- from warp.tests.test_volume_write import TestVolumeWrite
200
+ from warp.tests.tile.test_tile import TestTile
201
+ from warp.tests.tile.test_tile_mathdx import TestTileMathDx
202
+ from warp.tests.tile.test_tile_reduce import TestTileReduce
203
+ from warp.tests.tile.test_tile_shared_memory import TestTileSharedMemory
206
204
 
207
205
  test_classes = [
208
206
  TestAdam,
@@ -260,8 +258,6 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
260
258
  TestMatLite,
261
259
  TestMatScalarOps,
262
260
  TestMath,
263
- TestMatmul,
264
- TestMatmulLite,
265
261
  TestMempool,
266
262
  TestMesh,
267
263
  TestMeshQueryAABBMethods,
@@ -305,7 +301,7 @@ def default_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader)
305
301
  TestTriangleClosestPoint,
306
302
  TestTypes,
307
303
  TestUtils,
308
- TestVBD,
304
+ TestVbd,
309
305
  TestVec,
310
306
  TestVecLite,
311
307
  TestVecScalarOps,
@@ -322,65 +318,59 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
322
318
 
323
319
  Requires manual updates with test_ext.py for now.
324
320
  """
325
- from warp.tests.test_array import TestArray
321
+ from warp.tests.cuda.test_mempool import TestMempool
322
+ from warp.tests.cuda.test_peer import TestPeer
323
+ from warp.tests.cuda.test_pinned import TestPinned
324
+ from warp.tests.cuda.test_streams import TestStreams
325
+ from warp.tests.geometry.test_bvh import TestBvh
326
+ from warp.tests.geometry.test_hash_grid import TestHashGrid
327
+ from warp.tests.geometry.test_marching_cubes import TestMarchingCubes
328
+ from warp.tests.geometry.test_mesh import TestMesh
329
+ from warp.tests.geometry.test_mesh_query_aabb import TestMeshQueryAABBMethods
330
+ from warp.tests.geometry.test_mesh_query_point import TestMeshQueryPoint
331
+ from warp.tests.geometry.test_mesh_query_ray import TestMeshQueryRay
332
+ from warp.tests.geometry.test_volume import TestVolume
333
+ from warp.tests.geometry.test_volume_write import TestVolumeWrite
334
+ from warp.tests.interop.test_dlpack import TestDLPack
326
335
  from warp.tests.test_array_reduce import TestArrayReduce
327
336
  from warp.tests.test_bool import TestBool
328
337
  from warp.tests.test_builtins_resolution import TestBuiltinsResolution
329
- from warp.tests.test_bvh import TestBvh
330
338
  from warp.tests.test_codegen import TestCodeGen
331
339
  from warp.tests.test_compile_consts import TestConstants
332
340
  from warp.tests.test_conditional import TestConditional
333
341
  from warp.tests.test_copy import TestCopy
334
342
  from warp.tests.test_ctypes import TestCTypes
335
343
  from warp.tests.test_devices import TestDevices
336
- from warp.tests.test_dlpack import TestDLPack
337
- from warp.tests.test_fabricarray import TestFabricArray
338
344
  from warp.tests.test_fp16 import TestFp16
339
345
  from warp.tests.test_func import TestFunc
340
346
  from warp.tests.test_generics import TestGenerics
341
347
  from warp.tests.test_grad_customs import TestGradCustoms
342
348
  from warp.tests.test_grad_debug import TestGradDebug
343
- from warp.tests.test_hash_grid import TestHashGrid
344
349
  from warp.tests.test_indexedarray import TestIndexedArray
345
350
  from warp.tests.test_launch import TestLaunch
346
351
  from warp.tests.test_lvalue import TestLValue
347
- from warp.tests.test_marching_cubes import TestMarchingCubes
348
352
  from warp.tests.test_mat_lite import TestMatLite
349
353
  from warp.tests.test_math import TestMath
350
- from warp.tests.test_matmul_lite import TestMatmulLite
351
- from warp.tests.test_mempool import TestMempool
352
- from warp.tests.test_mesh import TestMesh
353
- from warp.tests.test_mesh_query_aabb import TestMeshQueryAABBMethods
354
- from warp.tests.test_mesh_query_point import TestMeshQueryPoint
355
- from warp.tests.test_mesh_query_ray import TestMeshQueryRay
356
354
  from warp.tests.test_mlp import TestMLP
357
355
  from warp.tests.test_module_hashing import TestModuleHashing
358
356
  from warp.tests.test_modules_lite import TestModuleLite
359
357
  from warp.tests.test_noise import TestNoise
360
358
  from warp.tests.test_operators import TestOperators
361
- from warp.tests.test_peer import TestPeer
362
- from warp.tests.test_pinned import TestPinned
363
- from warp.tests.test_quat import TestQuat
364
359
  from warp.tests.test_rand import TestRand
365
360
  from warp.tests.test_reload import TestReload
366
361
  from warp.tests.test_rounding import TestRounding
367
362
  from warp.tests.test_runlength_encode import TestRunlengthEncode
368
363
  from warp.tests.test_scalar_ops import TestScalarOps
369
364
  from warp.tests.test_snippet import TestSnippets
370
- from warp.tests.test_sparse import TestSparse
371
365
  from warp.tests.test_static import TestStatic
372
- from warp.tests.test_streams import TestStreams
373
366
  from warp.tests.test_tape import TestTape
374
- from warp.tests.test_tile_reduce import TestTileReduce
375
367
  from warp.tests.test_transient_module import TestTransientModule
376
368
  from warp.tests.test_types import TestTypes
377
369
  from warp.tests.test_utils import TestUtils
378
370
  from warp.tests.test_vec_lite import TestVecLite
379
- from warp.tests.test_volume import TestVolume
380
- from warp.tests.test_volume_write import TestVolumeWrite
371
+ from warp.tests.tile.test_tile_reduce import TestTileReduce
381
372
 
382
373
  test_classes = [
383
- TestArray,
384
374
  TestArrayReduce,
385
375
  TestBool,
386
376
  TestBuiltinsResolution,
@@ -392,7 +382,6 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
392
382
  TestCTypes,
393
383
  TestDevices,
394
384
  TestDLPack,
395
- TestFabricArray,
396
385
  TestFp16,
397
386
  TestFunc,
398
387
  TestGenerics,
@@ -405,7 +394,6 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
405
394
  TestMarchingCubes,
406
395
  TestMatLite,
407
396
  TestMath,
408
- TestMatmulLite,
409
397
  TestMempool,
410
398
  TestMesh,
411
399
  TestMeshQueryAABBMethods,
@@ -418,14 +406,12 @@ def kit_suite(test_loader: unittest.TestLoader = unittest.defaultTestLoader):
418
406
  TestOperators,
419
407
  TestPeer,
420
408
  TestPinned,
421
- TestQuat,
422
409
  TestRand,
423
410
  TestReload,
424
411
  TestRounding,
425
412
  TestRunlengthEncode,
426
413
  TestScalarOps,
427
414
  TestSnippets,
428
- TestSparse,
429
415
  TestStatic,
430
416
  TestStreams,
431
417
  TestTape,
@@ -392,7 +392,8 @@ def write_junit_results(
392
392
  error.text = str(test_data[4]) # Stacktrace
393
393
  elif test_status == "SKIP":
394
394
  skip = ET.SubElement(test_case, "skipped")
395
- skip.text = str(test_data[3]) # The skip reason
395
+ # Set the skip reason
396
+ skip.set("message", str(test_data[3]))
396
397
 
397
398
  tree = ET.ElementTree(root)
398
399
 
@@ -178,6 +178,7 @@ def main(argv=None):
178
178
  import warp as wp # NVIDIA Modification
179
179
 
180
180
  # Clear the Warp cache (NVIDIA Modification)
181
+ wp.clear_lto_cache()
181
182
  wp.clear_kernel_cache()
182
183
  print("Cleared Warp kernel cache")
183
184
 
@@ -558,7 +559,8 @@ def initialize_test_process(lock, shared_index, args, temp_dir):
558
559
 
559
560
  wp.config.kernel_cache_dir = cache_root_dir
560
561
 
561
- wp.build.clear_kernel_cache()
562
+ wp.clear_lto_cache()
563
+ wp.clear_kernel_cache()
562
564
  elif "WARP_CACHE_ROOT" in os.environ:
563
565
  # Using a shared cache for all test processes
564
566
  wp.config.kernel_cache_dir = os.path.join(os.getenv("WARP_CACHE_ROOT"), wp.config.version)