warp-lang 1.0.2__py3-none-win_amd64.whl → 1.2.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.
- warp/__init__.py +108 -97
- warp/__init__.pyi +1 -1
- warp/bin/warp-clang.dll +0 -0
- warp/bin/warp.dll +0 -0
- warp/build.py +88 -113
- warp/build_dll.py +383 -375
- warp/builtins.py +3693 -3354
- warp/codegen.py +2925 -2792
- warp/config.py +40 -36
- warp/constants.py +49 -45
- warp/context.py +5409 -5102
- warp/dlpack.py +442 -442
- warp/examples/__init__.py +16 -16
- warp/examples/assets/bear.usd +0 -0
- warp/examples/assets/bunny.usd +0 -0
- warp/examples/assets/cartpole.urdf +110 -110
- warp/examples/assets/crazyflie.usd +0 -0
- warp/examples/assets/cube.usd +0 -0
- warp/examples/assets/nv_ant.xml +92 -92
- warp/examples/assets/nv_humanoid.xml +183 -183
- warp/examples/assets/quadruped.urdf +267 -267
- warp/examples/assets/rocks.nvdb +0 -0
- warp/examples/assets/rocks.usd +0 -0
- warp/examples/assets/sphere.usd +0 -0
- warp/examples/benchmarks/benchmark_api.py +381 -383
- warp/examples/benchmarks/benchmark_cloth.py +278 -277
- warp/examples/benchmarks/benchmark_cloth_cupy.py +88 -88
- warp/examples/benchmarks/benchmark_cloth_jax.py +97 -100
- warp/examples/benchmarks/benchmark_cloth_numba.py +146 -142
- warp/examples/benchmarks/benchmark_cloth_numpy.py +77 -77
- warp/examples/benchmarks/benchmark_cloth_pytorch.py +86 -86
- warp/examples/benchmarks/benchmark_cloth_taichi.py +112 -112
- warp/examples/benchmarks/benchmark_cloth_warp.py +145 -146
- warp/examples/benchmarks/benchmark_launches.py +293 -295
- warp/examples/browse.py +29 -29
- warp/examples/core/example_dem.py +232 -219
- warp/examples/core/example_fluid.py +291 -267
- warp/examples/core/example_graph_capture.py +142 -126
- warp/examples/core/example_marching_cubes.py +186 -174
- warp/examples/core/example_mesh.py +172 -155
- warp/examples/core/example_mesh_intersect.py +203 -193
- warp/examples/core/example_nvdb.py +174 -170
- warp/examples/core/example_raycast.py +103 -90
- warp/examples/core/example_raymarch.py +197 -178
- warp/examples/core/example_render_opengl.py +183 -141
- warp/examples/core/example_sph.py +403 -387
- warp/examples/core/example_torch.py +219 -181
- warp/examples/core/example_wave.py +261 -248
- warp/examples/fem/bsr_utils.py +378 -380
- warp/examples/fem/example_apic_fluid.py +432 -389
- warp/examples/fem/example_burgers.py +262 -0
- warp/examples/fem/example_convection_diffusion.py +180 -168
- warp/examples/fem/example_convection_diffusion_dg.py +217 -209
- warp/examples/fem/example_deformed_geometry.py +175 -159
- warp/examples/fem/example_diffusion.py +199 -173
- warp/examples/fem/example_diffusion_3d.py +178 -152
- warp/examples/fem/example_diffusion_mgpu.py +219 -214
- warp/examples/fem/example_mixed_elasticity.py +242 -222
- warp/examples/fem/example_navier_stokes.py +257 -243
- warp/examples/fem/example_stokes.py +218 -192
- warp/examples/fem/example_stokes_transfer.py +263 -249
- warp/examples/fem/mesh_utils.py +133 -109
- warp/examples/fem/plot_utils.py +292 -287
- warp/examples/optim/example_bounce.py +258 -246
- warp/examples/optim/example_cloth_throw.py +220 -209
- warp/examples/optim/example_diffray.py +564 -536
- warp/examples/optim/example_drone.py +862 -835
- warp/examples/optim/example_inverse_kinematics.py +174 -168
- warp/examples/optim/example_inverse_kinematics_torch.py +183 -169
- warp/examples/optim/example_spring_cage.py +237 -231
- warp/examples/optim/example_trajectory.py +221 -199
- warp/examples/optim/example_walker.py +304 -293
- warp/examples/sim/example_cartpole.py +137 -129
- warp/examples/sim/example_cloth.py +194 -186
- warp/examples/sim/example_granular.py +122 -111
- warp/examples/sim/example_granular_collision_sdf.py +195 -186
- warp/examples/sim/example_jacobian_ik.py +234 -214
- warp/examples/sim/example_particle_chain.py +116 -105
- warp/examples/sim/example_quadruped.py +191 -180
- warp/examples/sim/example_rigid_chain.py +195 -187
- warp/examples/sim/example_rigid_contact.py +187 -177
- warp/examples/sim/example_rigid_force.py +125 -125
- warp/examples/sim/example_rigid_gyroscopic.py +107 -95
- warp/examples/sim/example_rigid_soft_contact.py +132 -122
- warp/examples/sim/example_soft_body.py +188 -177
- warp/fabric.py +337 -335
- warp/fem/__init__.py +61 -27
- warp/fem/cache.py +403 -388
- warp/fem/dirichlet.py +178 -179
- warp/fem/domain.py +262 -263
- warp/fem/field/__init__.py +100 -101
- warp/fem/field/field.py +148 -149
- warp/fem/field/nodal_field.py +298 -299
- warp/fem/field/restriction.py +22 -21
- warp/fem/field/test.py +180 -181
- warp/fem/field/trial.py +183 -183
- warp/fem/geometry/__init__.py +16 -19
- warp/fem/geometry/closest_point.py +69 -70
- warp/fem/geometry/deformed_geometry.py +270 -271
- warp/fem/geometry/element.py +748 -744
- warp/fem/geometry/geometry.py +184 -186
- warp/fem/geometry/grid_2d.py +380 -373
- warp/fem/geometry/grid_3d.py +437 -435
- warp/fem/geometry/hexmesh.py +953 -953
- warp/fem/geometry/nanogrid.py +455 -0
- warp/fem/geometry/partition.py +374 -376
- warp/fem/geometry/quadmesh_2d.py +532 -532
- warp/fem/geometry/tetmesh.py +840 -840
- warp/fem/geometry/trimesh_2d.py +577 -577
- warp/fem/integrate.py +1684 -1615
- warp/fem/operator.py +190 -191
- warp/fem/polynomial.py +214 -213
- warp/fem/quadrature/__init__.py +2 -2
- warp/fem/quadrature/pic_quadrature.py +243 -245
- warp/fem/quadrature/quadrature.py +295 -294
- warp/fem/space/__init__.py +179 -292
- warp/fem/space/basis_space.py +522 -489
- warp/fem/space/collocated_function_space.py +100 -105
- warp/fem/space/dof_mapper.py +236 -236
- warp/fem/space/function_space.py +148 -145
- warp/fem/space/grid_2d_function_space.py +148 -267
- warp/fem/space/grid_3d_function_space.py +167 -306
- warp/fem/space/hexmesh_function_space.py +253 -352
- warp/fem/space/nanogrid_function_space.py +202 -0
- warp/fem/space/partition.py +350 -350
- warp/fem/space/quadmesh_2d_function_space.py +261 -369
- warp/fem/space/restriction.py +161 -160
- warp/fem/space/shape/__init__.py +90 -15
- warp/fem/space/shape/cube_shape_function.py +728 -738
- warp/fem/space/shape/shape_function.py +102 -103
- warp/fem/space/shape/square_shape_function.py +611 -611
- warp/fem/space/shape/tet_shape_function.py +565 -567
- warp/fem/space/shape/triangle_shape_function.py +429 -429
- warp/fem/space/tetmesh_function_space.py +224 -292
- warp/fem/space/topology.py +297 -295
- warp/fem/space/trimesh_2d_function_space.py +153 -221
- warp/fem/types.py +77 -77
- warp/fem/utils.py +495 -495
- warp/jax.py +166 -141
- warp/jax_experimental.py +341 -339
- warp/native/array.h +1081 -1025
- warp/native/builtin.h +1603 -1560
- warp/native/bvh.cpp +402 -398
- warp/native/bvh.cu +533 -525
- warp/native/bvh.h +430 -429
- warp/native/clang/clang.cpp +496 -464
- warp/native/crt.cpp +42 -32
- warp/native/crt.h +352 -335
- warp/native/cuda_crt.h +1049 -1049
- warp/native/cuda_util.cpp +549 -540
- warp/native/cuda_util.h +288 -203
- warp/native/cutlass_gemm.cpp +34 -34
- warp/native/cutlass_gemm.cu +372 -372
- warp/native/error.cpp +66 -66
- warp/native/error.h +27 -27
- warp/native/exports.h +187 -0
- warp/native/fabric.h +228 -228
- warp/native/hashgrid.cpp +301 -278
- warp/native/hashgrid.cu +78 -77
- warp/native/hashgrid.h +227 -227
- warp/native/initializer_array.h +32 -32
- warp/native/intersect.h +1204 -1204
- warp/native/intersect_adj.h +365 -365
- warp/native/intersect_tri.h +322 -322
- warp/native/marching.cpp +2 -2
- warp/native/marching.cu +497 -497
- warp/native/marching.h +2 -2
- warp/native/mat.h +1545 -1498
- warp/native/matnn.h +333 -333
- warp/native/mesh.cpp +203 -203
- warp/native/mesh.cu +292 -293
- warp/native/mesh.h +1887 -1887
- warp/native/nanovdb/GridHandle.h +366 -0
- warp/native/nanovdb/HostBuffer.h +590 -0
- warp/native/nanovdb/NanoVDB.h +6624 -4782
- warp/native/nanovdb/PNanoVDB.h +3390 -2553
- warp/native/noise.h +850 -850
- warp/native/quat.h +1112 -1085
- warp/native/rand.h +303 -299
- warp/native/range.h +108 -108
- warp/native/reduce.cpp +156 -156
- warp/native/reduce.cu +348 -348
- warp/native/runlength_encode.cpp +61 -61
- warp/native/runlength_encode.cu +46 -46
- warp/native/scan.cpp +30 -30
- warp/native/scan.cu +36 -36
- warp/native/scan.h +7 -7
- warp/native/solid_angle.h +442 -442
- warp/native/sort.cpp +94 -94
- warp/native/sort.cu +97 -97
- warp/native/sort.h +14 -14
- warp/native/sparse.cpp +337 -337
- warp/native/sparse.cu +544 -544
- warp/native/spatial.h +630 -630
- warp/native/svd.h +562 -562
- warp/native/temp_buffer.h +30 -30
- warp/native/vec.h +1177 -1133
- warp/native/volume.cpp +529 -297
- warp/native/volume.cu +58 -32
- warp/native/volume.h +960 -538
- warp/native/volume_builder.cu +446 -425
- warp/native/volume_builder.h +34 -19
- warp/native/volume_impl.h +61 -0
- warp/native/warp.cpp +1057 -1052
- warp/native/warp.cu +2949 -2828
- warp/native/warp.h +321 -305
- warp/optim/__init__.py +9 -9
- warp/optim/adam.py +120 -120
- warp/optim/linear.py +1104 -939
- warp/optim/sgd.py +104 -92
- warp/render/__init__.py +10 -10
- warp/render/render_opengl.py +3356 -3204
- warp/render/render_usd.py +768 -749
- warp/render/utils.py +152 -150
- warp/sim/__init__.py +52 -59
- warp/sim/articulation.py +685 -685
- warp/sim/collide.py +1594 -1590
- warp/sim/import_mjcf.py +489 -481
- warp/sim/import_snu.py +220 -221
- warp/sim/import_urdf.py +536 -516
- warp/sim/import_usd.py +887 -881
- warp/sim/inertia.py +316 -317
- warp/sim/integrator.py +234 -233
- warp/sim/integrator_euler.py +1956 -1956
- warp/sim/integrator_featherstone.py +1917 -1991
- warp/sim/integrator_xpbd.py +3288 -3312
- warp/sim/model.py +4473 -4314
- warp/sim/particles.py +113 -112
- warp/sim/render.py +417 -403
- warp/sim/utils.py +413 -410
- warp/sparse.py +1289 -1227
- warp/stubs.py +2192 -2469
- warp/tape.py +1162 -225
- warp/tests/__init__.py +1 -1
- warp/tests/__main__.py +4 -4
- warp/tests/assets/test_index_grid.nvdb +0 -0
- warp/tests/assets/torus.usda +105 -105
- warp/tests/aux_test_class_kernel.py +26 -26
- warp/tests/aux_test_compile_consts_dummy.py +10 -10
- warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -21
- warp/tests/aux_test_dependent.py +20 -22
- warp/tests/aux_test_grad_customs.py +21 -23
- warp/tests/aux_test_reference.py +9 -11
- warp/tests/aux_test_reference_reference.py +8 -10
- warp/tests/aux_test_square.py +15 -17
- warp/tests/aux_test_unresolved_func.py +14 -14
- warp/tests/aux_test_unresolved_symbol.py +14 -14
- warp/tests/disabled_kinematics.py +237 -239
- warp/tests/run_coverage_serial.py +31 -31
- warp/tests/test_adam.py +155 -157
- warp/tests/test_arithmetic.py +1088 -1124
- warp/tests/test_array.py +2415 -2326
- warp/tests/test_array_reduce.py +148 -150
- warp/tests/test_async.py +666 -656
- warp/tests/test_atomic.py +139 -141
- warp/tests/test_bool.py +212 -149
- warp/tests/test_builtins_resolution.py +1290 -1292
- warp/tests/test_bvh.py +162 -171
- warp/tests/test_closest_point_edge_edge.py +227 -228
- warp/tests/test_codegen.py +562 -553
- warp/tests/test_compile_consts.py +217 -101
- warp/tests/test_conditional.py +244 -246
- warp/tests/test_copy.py +230 -215
- warp/tests/test_ctypes.py +630 -632
- warp/tests/test_dense.py +65 -67
- warp/tests/test_devices.py +89 -98
- warp/tests/test_dlpack.py +528 -529
- warp/tests/test_examples.py +403 -378
- warp/tests/test_fabricarray.py +952 -955
- warp/tests/test_fast_math.py +60 -54
- warp/tests/test_fem.py +1298 -1278
- warp/tests/test_fp16.py +128 -130
- warp/tests/test_func.py +336 -337
- warp/tests/test_generics.py +596 -571
- warp/tests/test_grad.py +885 -640
- warp/tests/test_grad_customs.py +331 -336
- warp/tests/test_hash_grid.py +208 -164
- warp/tests/test_import.py +37 -39
- warp/tests/test_indexedarray.py +1132 -1134
- warp/tests/test_intersect.py +65 -67
- warp/tests/test_jax.py +305 -307
- warp/tests/test_large.py +169 -164
- warp/tests/test_launch.py +352 -354
- warp/tests/test_lerp.py +217 -261
- warp/tests/test_linear_solvers.py +189 -171
- warp/tests/test_lvalue.py +419 -493
- warp/tests/test_marching_cubes.py +63 -65
- warp/tests/test_mat.py +1799 -1827
- warp/tests/test_mat_lite.py +113 -115
- warp/tests/test_mat_scalar_ops.py +2905 -2889
- warp/tests/test_math.py +124 -193
- warp/tests/test_matmul.py +498 -499
- warp/tests/test_matmul_lite.py +408 -410
- warp/tests/test_mempool.py +186 -190
- warp/tests/test_mesh.py +281 -324
- warp/tests/test_mesh_query_aabb.py +226 -241
- warp/tests/test_mesh_query_point.py +690 -702
- warp/tests/test_mesh_query_ray.py +290 -303
- warp/tests/test_mlp.py +274 -276
- warp/tests/test_model.py +108 -110
- warp/tests/test_module_hashing.py +111 -0
- warp/tests/test_modules_lite.py +36 -39
- warp/tests/test_multigpu.py +161 -163
- warp/tests/test_noise.py +244 -248
- warp/tests/test_operators.py +248 -250
- warp/tests/test_options.py +121 -125
- warp/tests/test_peer.py +131 -137
- warp/tests/test_pinned.py +76 -78
- warp/tests/test_print.py +52 -54
- warp/tests/test_quat.py +2084 -2086
- warp/tests/test_rand.py +324 -288
- warp/tests/test_reload.py +207 -217
- warp/tests/test_rounding.py +177 -179
- warp/tests/test_runlength_encode.py +188 -190
- warp/tests/test_sim_grad.py +241 -0
- warp/tests/test_sim_kinematics.py +89 -97
- warp/tests/test_smoothstep.py +166 -168
- warp/tests/test_snippet.py +303 -266
- warp/tests/test_sparse.py +466 -460
- warp/tests/test_spatial.py +2146 -2148
- warp/tests/test_special_values.py +362 -0
- warp/tests/test_streams.py +484 -473
- warp/tests/test_struct.py +708 -675
- warp/tests/test_tape.py +171 -148
- warp/tests/test_torch.py +741 -743
- warp/tests/test_transient_module.py +85 -87
- warp/tests/test_types.py +554 -659
- warp/tests/test_utils.py +488 -499
- warp/tests/test_vec.py +1262 -1268
- warp/tests/test_vec_lite.py +71 -73
- warp/tests/test_vec_scalar_ops.py +2097 -2099
- warp/tests/test_verify_fp.py +92 -94
- warp/tests/test_volume.py +961 -736
- warp/tests/test_volume_write.py +338 -265
- warp/tests/unittest_serial.py +38 -37
- warp/tests/unittest_suites.py +367 -359
- warp/tests/unittest_utils.py +434 -578
- warp/tests/unused_test_misc.py +69 -71
- warp/tests/walkthrough_debug.py +85 -85
- warp/thirdparty/appdirs.py +598 -598
- warp/thirdparty/dlpack.py +143 -143
- warp/thirdparty/unittest_parallel.py +563 -561
- warp/torch.py +321 -295
- warp/types.py +4941 -4450
- warp/utils.py +1008 -821
- {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/LICENSE.md +126 -126
- {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/METADATA +365 -400
- warp_lang-1.2.0.dist-info/RECORD +359 -0
- warp/examples/assets/cube.usda +0 -42
- warp/examples/assets/sphere.usda +0 -56
- warp/examples/assets/torus.usda +0 -105
- warp/examples/fem/example_convection_diffusion_dg0.py +0 -194
- warp/native/nanovdb/PNanoVDBWrite.h +0 -295
- warp_lang-1.0.2.dist-info/RECORD +0 -352
- {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/WHEEL +0 -0
- {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/top_level.txt +0 -0
warp/tests/test_generics.py
CHANGED
|
@@ -1,571 +1,596 @@
|
|
|
1
|
-
# Copyright (c) 2023 NVIDIA CORPORATION. All rights reserved.
|
|
2
|
-
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
3
|
-
# and proprietary rights in and to this software, related documentation
|
|
4
|
-
# and any modifications thereto. Any use, reproduction, disclosure or
|
|
5
|
-
# distribution of this software and related documentation without an express
|
|
6
|
-
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
7
|
-
|
|
8
|
-
import unittest
|
|
9
|
-
from typing import Any
|
|
10
|
-
|
|
11
|
-
import numpy as np
|
|
12
|
-
|
|
13
|
-
import warp as wp
|
|
14
|
-
from warp.tests.unittest_utils import *
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
wp.
|
|
27
|
-
wp.
|
|
28
|
-
v1
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
wp.expect_eq(specialized_func(
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
wp.
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
v2
|
|
77
|
-
v3
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
tid =
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
wp.overload(generic_array_kernel_v1, [wp.array(dtype=
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
tid =
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
wp.overload(generic_array_kernel_v2, [wp.array(dtype=
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
tid =
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
wp.overload(generic_array_kernel_v3, [wp.array(dtype=
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
n =
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
wp.launch(generic_array_kernel_v1, dim=n, inputs=[
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
assert_np_equal(
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
wp.launch(generic_array_kernel_v2, dim=n, inputs=[
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
assert_np_equal(
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
wp.launch(generic_array_kernel_v3, dim=n, inputs=[
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
assert_np_equal(
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
wp.overload(generic_accumulator_kernel, {"value":
|
|
165
|
-
wp.overload(generic_accumulator_kernel, {"value":
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a,
|
|
176
|
-
wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a,
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
n =
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
wp.launch(generic_fill, dim=
|
|
199
|
-
wp.launch(generic_fill, dim=
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
assert_np_equal(
|
|
204
|
-
assert_np_equal(
|
|
205
|
-
|
|
206
|
-
|
|
207
|
-
|
|
208
|
-
|
|
209
|
-
|
|
210
|
-
|
|
211
|
-
|
|
212
|
-
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
n =
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
232
|
-
wp.launch(
|
|
233
|
-
wp.launch(
|
|
234
|
-
|
|
235
|
-
|
|
236
|
-
|
|
237
|
-
assert_np_equal(
|
|
238
|
-
assert_np_equal(
|
|
239
|
-
|
|
240
|
-
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
|
|
254
|
-
|
|
255
|
-
|
|
256
|
-
|
|
257
|
-
|
|
258
|
-
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
264
|
-
|
|
265
|
-
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
278
|
-
|
|
279
|
-
|
|
280
|
-
|
|
281
|
-
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
|
|
285
|
-
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
|
|
301
|
-
|
|
302
|
-
|
|
303
|
-
|
|
304
|
-
|
|
305
|
-
|
|
306
|
-
|
|
307
|
-
|
|
308
|
-
|
|
309
|
-
|
|
310
|
-
wp.
|
|
311
|
-
|
|
312
|
-
|
|
313
|
-
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
320
|
-
|
|
321
|
-
|
|
322
|
-
|
|
323
|
-
|
|
324
|
-
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
338
|
-
|
|
339
|
-
|
|
340
|
-
|
|
341
|
-
|
|
342
|
-
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
|
|
348
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
352
|
-
|
|
353
|
-
|
|
354
|
-
|
|
355
|
-
|
|
356
|
-
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
|
|
364
|
-
|
|
365
|
-
|
|
366
|
-
|
|
367
|
-
|
|
368
|
-
|
|
369
|
-
|
|
370
|
-
|
|
371
|
-
|
|
372
|
-
|
|
373
|
-
|
|
374
|
-
|
|
375
|
-
|
|
376
|
-
|
|
377
|
-
@wp.
|
|
378
|
-
|
|
379
|
-
|
|
380
|
-
|
|
381
|
-
|
|
382
|
-
|
|
383
|
-
|
|
384
|
-
wp.
|
|
385
|
-
|
|
386
|
-
|
|
387
|
-
|
|
388
|
-
|
|
389
|
-
|
|
390
|
-
|
|
391
|
-
|
|
392
|
-
|
|
393
|
-
|
|
394
|
-
|
|
395
|
-
|
|
396
|
-
|
|
397
|
-
|
|
398
|
-
|
|
399
|
-
|
|
400
|
-
|
|
401
|
-
|
|
402
|
-
|
|
403
|
-
|
|
404
|
-
|
|
405
|
-
wp.expect_eq(
|
|
406
|
-
|
|
407
|
-
|
|
408
|
-
wp.overload(
|
|
409
|
-
wp.overload(
|
|
410
|
-
|
|
411
|
-
|
|
412
|
-
|
|
413
|
-
|
|
414
|
-
|
|
415
|
-
|
|
416
|
-
|
|
417
|
-
wp.
|
|
418
|
-
|
|
419
|
-
|
|
420
|
-
|
|
421
|
-
|
|
422
|
-
|
|
423
|
-
|
|
424
|
-
|
|
425
|
-
|
|
426
|
-
|
|
427
|
-
|
|
428
|
-
|
|
429
|
-
wp.
|
|
430
|
-
|
|
431
|
-
|
|
432
|
-
|
|
433
|
-
|
|
434
|
-
|
|
435
|
-
|
|
436
|
-
|
|
437
|
-
|
|
438
|
-
|
|
439
|
-
|
|
440
|
-
|
|
441
|
-
|
|
442
|
-
|
|
443
|
-
|
|
444
|
-
|
|
445
|
-
|
|
446
|
-
|
|
447
|
-
|
|
448
|
-
|
|
449
|
-
wp.
|
|
450
|
-
|
|
451
|
-
|
|
452
|
-
|
|
453
|
-
|
|
454
|
-
|
|
455
|
-
|
|
456
|
-
|
|
457
|
-
|
|
458
|
-
|
|
459
|
-
wp.
|
|
460
|
-
|
|
461
|
-
|
|
462
|
-
|
|
463
|
-
|
|
464
|
-
|
|
465
|
-
|
|
466
|
-
|
|
467
|
-
|
|
468
|
-
|
|
469
|
-
|
|
470
|
-
|
|
471
|
-
|
|
472
|
-
|
|
473
|
-
|
|
474
|
-
|
|
475
|
-
|
|
476
|
-
|
|
477
|
-
|
|
478
|
-
|
|
479
|
-
|
|
480
|
-
|
|
481
|
-
|
|
482
|
-
|
|
483
|
-
|
|
484
|
-
|
|
485
|
-
|
|
486
|
-
|
|
487
|
-
|
|
488
|
-
|
|
489
|
-
|
|
490
|
-
|
|
491
|
-
|
|
492
|
-
|
|
493
|
-
|
|
494
|
-
|
|
495
|
-
|
|
496
|
-
|
|
497
|
-
|
|
498
|
-
)
|
|
499
|
-
|
|
500
|
-
|
|
501
|
-
|
|
502
|
-
|
|
503
|
-
|
|
504
|
-
|
|
505
|
-
|
|
506
|
-
|
|
507
|
-
|
|
508
|
-
|
|
509
|
-
|
|
510
|
-
|
|
511
|
-
|
|
512
|
-
|
|
513
|
-
|
|
514
|
-
|
|
515
|
-
|
|
516
|
-
|
|
517
|
-
|
|
518
|
-
|
|
519
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
523
|
-
|
|
524
|
-
|
|
525
|
-
|
|
526
|
-
|
|
527
|
-
|
|
528
|
-
|
|
529
|
-
|
|
530
|
-
|
|
531
|
-
|
|
532
|
-
|
|
533
|
-
|
|
534
|
-
|
|
535
|
-
|
|
536
|
-
|
|
537
|
-
|
|
538
|
-
|
|
539
|
-
|
|
540
|
-
|
|
541
|
-
|
|
542
|
-
|
|
543
|
-
|
|
544
|
-
|
|
545
|
-
|
|
546
|
-
|
|
547
|
-
|
|
548
|
-
)
|
|
549
|
-
|
|
550
|
-
|
|
551
|
-
|
|
552
|
-
|
|
553
|
-
|
|
554
|
-
|
|
555
|
-
|
|
556
|
-
|
|
557
|
-
|
|
558
|
-
add_kernel_test(
|
|
559
|
-
TestGenerics,
|
|
560
|
-
name="
|
|
561
|
-
kernel=
|
|
562
|
-
dim=1,
|
|
563
|
-
inputs=[
|
|
564
|
-
devices=devices,
|
|
565
|
-
)
|
|
566
|
-
|
|
567
|
-
|
|
568
|
-
|
|
569
|
-
|
|
570
|
-
|
|
571
|
-
|
|
1
|
+
# Copyright (c) 2023 NVIDIA CORPORATION. All rights reserved.
|
|
2
|
+
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
3
|
+
# and proprietary rights in and to this software, related documentation
|
|
4
|
+
# and any modifications thereto. Any use, reproduction, disclosure or
|
|
5
|
+
# distribution of this software and related documentation without an express
|
|
6
|
+
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
7
|
+
|
|
8
|
+
import unittest
|
|
9
|
+
from typing import Any
|
|
10
|
+
|
|
11
|
+
import numpy as np
|
|
12
|
+
|
|
13
|
+
import warp as wp
|
|
14
|
+
from warp.tests.unittest_utils import *
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
@wp.func
|
|
18
|
+
def generic_adder(a: Any, b: Any):
|
|
19
|
+
return a + b
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
@wp.kernel
|
|
23
|
+
def test_generic_adder():
|
|
24
|
+
wp.expect_eq(generic_adder(17, 25), 42)
|
|
25
|
+
wp.expect_eq(generic_adder(7.0, 10.0), 17.0)
|
|
26
|
+
v1 = wp.vec3(1.0, 2.0, 3.0)
|
|
27
|
+
v2 = wp.vec3(10.0, 20.0, 30.0)
|
|
28
|
+
wp.expect_eq(generic_adder(v1, v2), wp.vec3(11.0, 22.0, 33.0))
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
# regular functions for floats
|
|
32
|
+
@wp.func
|
|
33
|
+
def specialized_func(a: float, b: float):
|
|
34
|
+
return a * b
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@wp.func
|
|
38
|
+
def specialized_func(a: float, b: float, c: float):
|
|
39
|
+
return a * b * c
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
# generic forms
|
|
43
|
+
@wp.func
|
|
44
|
+
def specialized_func(a: Any, b: Any):
|
|
45
|
+
return a + b
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
@wp.func
|
|
49
|
+
def specialized_func(a: Any, b: Any, c: Any):
|
|
50
|
+
return a + b + c
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
# specializations for ints
|
|
54
|
+
@wp.func
|
|
55
|
+
def specialized_func(a: int, b: int):
|
|
56
|
+
return a - b
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
@wp.func
|
|
60
|
+
def specialized_func(a: int, b: int, c: int):
|
|
61
|
+
return a - b - c
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
@wp.kernel
|
|
65
|
+
def test_specialized_func():
|
|
66
|
+
# subtraction with int args
|
|
67
|
+
wp.expect_eq(specialized_func(17, 25), -8)
|
|
68
|
+
wp.expect_eq(specialized_func(17, 25, 10), -18)
|
|
69
|
+
# multiplication with float args
|
|
70
|
+
wp.expect_eq(specialized_func(7.0, 10.0), 70.0)
|
|
71
|
+
wp.expect_eq(specialized_func(7.0, 10.0, 2.0), 140.0)
|
|
72
|
+
# addition with vector args
|
|
73
|
+
v1 = wp.vec3(1.0, 2.0, 3.0)
|
|
74
|
+
v2 = wp.vec3(10.0, 20.0, 30.0)
|
|
75
|
+
v3 = wp.vec3(100.0, 200.0, 300.0)
|
|
76
|
+
wp.expect_eq(specialized_func(v1, v2), wp.vec3(11.0, 22.0, 33.0))
|
|
77
|
+
wp.expect_eq(specialized_func(v1, v2, v3), wp.vec3(111.0, 222.0, 333.0))
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
# generic array kernel, version 1 (Any)
|
|
81
|
+
@wp.kernel
|
|
82
|
+
def generic_array_kernel_v1(a: Any, b: Any, c: Any):
|
|
83
|
+
tid = wp.tid()
|
|
84
|
+
sum = a[tid] + b[tid] # test direct access
|
|
85
|
+
c[tid] = generic_adder(sum, sum) # test generic function
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
wp.overload(generic_array_kernel_v1, [wp.array(dtype=int), wp.array(dtype=int), wp.array(dtype=int)])
|
|
89
|
+
wp.overload(generic_array_kernel_v1, [wp.array(dtype=float), wp.array(dtype=float), wp.array(dtype=float)])
|
|
90
|
+
wp.overload(generic_array_kernel_v1, [wp.array(dtype=wp.vec3), wp.array(dtype=wp.vec3), wp.array(dtype=wp.vec3)])
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
# generic array kernel, version 2 (generic dtype)
|
|
94
|
+
@wp.kernel
|
|
95
|
+
def generic_array_kernel_v2(a: wp.array(dtype=Any), b: wp.array(dtype=Any), c: wp.array(dtype=Any)):
|
|
96
|
+
tid = wp.tid()
|
|
97
|
+
sum = a[tid] + b[tid] # test direct access
|
|
98
|
+
c[tid] = generic_adder(sum, sum) # test generic function
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
wp.overload(generic_array_kernel_v2, [wp.array(dtype=int), wp.array(dtype=int), wp.array(dtype=int)])
|
|
102
|
+
wp.overload(generic_array_kernel_v2, [wp.array(dtype=float), wp.array(dtype=float), wp.array(dtype=float)])
|
|
103
|
+
wp.overload(generic_array_kernel_v2, [wp.array(dtype=wp.vec3), wp.array(dtype=wp.vec3), wp.array(dtype=wp.vec3)])
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
# generic array kernel, version 3 (unspecified dtype)
|
|
107
|
+
@wp.kernel
|
|
108
|
+
def generic_array_kernel_v3(a: wp.array(), b: wp.array(), c: wp.array()):
|
|
109
|
+
tid = wp.tid()
|
|
110
|
+
sum = a[tid] + b[tid] # test direct access
|
|
111
|
+
c[tid] = generic_adder(sum, sum) # test generic function
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
wp.overload(generic_array_kernel_v3, [wp.array(dtype=int), wp.array(dtype=int), wp.array(dtype=int)])
|
|
115
|
+
wp.overload(generic_array_kernel_v3, [wp.array(dtype=float), wp.array(dtype=float), wp.array(dtype=float)])
|
|
116
|
+
wp.overload(generic_array_kernel_v3, [wp.array(dtype=wp.vec3), wp.array(dtype=wp.vec3), wp.array(dtype=wp.vec3)])
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def test_generic_array_kernel(test, device):
|
|
120
|
+
with wp.ScopedDevice(device):
|
|
121
|
+
n = 10
|
|
122
|
+
|
|
123
|
+
ai = wp.array(data=np.ones(n, dtype=np.int32))
|
|
124
|
+
ci = wp.empty(10, dtype=int)
|
|
125
|
+
|
|
126
|
+
af = wp.array(data=np.ones(n, dtype=np.float32))
|
|
127
|
+
cf = wp.empty(10, dtype=float)
|
|
128
|
+
|
|
129
|
+
a3 = wp.array(data=np.ones((n, 3), dtype=np.float32), dtype=wp.vec3)
|
|
130
|
+
c3 = wp.empty(n, dtype=wp.vec3)
|
|
131
|
+
|
|
132
|
+
wp.launch(generic_array_kernel_v1, dim=n, inputs=[af, af, cf])
|
|
133
|
+
wp.launch(generic_array_kernel_v1, dim=n, inputs=[ai, ai, ci])
|
|
134
|
+
wp.launch(generic_array_kernel_v1, dim=n, inputs=[a3, a3, c3])
|
|
135
|
+
assert_np_equal(ci.numpy(), np.full((n,), 4, dtype=np.int32))
|
|
136
|
+
assert_np_equal(cf.numpy(), np.full((n,), 4.0, dtype=np.float32))
|
|
137
|
+
assert_np_equal(c3.numpy(), np.full((n, 3), 4.0, dtype=np.float32))
|
|
138
|
+
|
|
139
|
+
wp.launch(generic_array_kernel_v2, dim=n, inputs=[af, af, cf])
|
|
140
|
+
wp.launch(generic_array_kernel_v2, dim=n, inputs=[ai, ai, ci])
|
|
141
|
+
wp.launch(generic_array_kernel_v2, dim=n, inputs=[a3, a3, c3])
|
|
142
|
+
assert_np_equal(ci.numpy(), np.full((n,), 4, dtype=np.int32))
|
|
143
|
+
assert_np_equal(cf.numpy(), np.full((n,), 4.0, dtype=np.float32))
|
|
144
|
+
assert_np_equal(c3.numpy(), np.full((n, 3), 4.0, dtype=np.float32))
|
|
145
|
+
|
|
146
|
+
wp.launch(generic_array_kernel_v3, dim=n, inputs=[af, af, cf])
|
|
147
|
+
wp.launch(generic_array_kernel_v3, dim=n, inputs=[ai, ai, ci])
|
|
148
|
+
wp.launch(generic_array_kernel_v3, dim=n, inputs=[a3, a3, c3])
|
|
149
|
+
assert_np_equal(ci.numpy(), np.full((n,), 4, dtype=np.int32))
|
|
150
|
+
assert_np_equal(cf.numpy(), np.full((n,), 4.0, dtype=np.float32))
|
|
151
|
+
assert_np_equal(c3.numpy(), np.full((n, 3), 4.0, dtype=np.float32))
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
# kernel that adds any scalar value to an array
|
|
155
|
+
@wp.kernel
|
|
156
|
+
def generic_accumulator_kernel(a: wp.array(dtype=wp.float64), value: Any):
|
|
157
|
+
tid = wp.tid()
|
|
158
|
+
a[tid] = a[tid] + wp.float64(value)
|
|
159
|
+
|
|
160
|
+
|
|
161
|
+
# overload named args
|
|
162
|
+
wp.overload(generic_accumulator_kernel, {"value": int})
|
|
163
|
+
wp.overload(generic_accumulator_kernel, {"value": float})
|
|
164
|
+
wp.overload(generic_accumulator_kernel, {"value": wp.float64})
|
|
165
|
+
wp.overload(generic_accumulator_kernel, {"value": wp.bool})
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
def test_generic_accumulator_kernel(test, device):
|
|
169
|
+
with wp.ScopedDevice(device):
|
|
170
|
+
n = 10
|
|
171
|
+
a = wp.zeros(n, dtype=wp.float64)
|
|
172
|
+
|
|
173
|
+
wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, 25])
|
|
174
|
+
wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, 17.0])
|
|
175
|
+
wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, wp.float64(8.0)])
|
|
176
|
+
wp.launch(generic_accumulator_kernel, dim=a.size, inputs=[a, wp.bool(True)])
|
|
177
|
+
|
|
178
|
+
assert_np_equal(a.numpy(), np.full((n,), 51.0, dtype=np.float64))
|
|
179
|
+
|
|
180
|
+
|
|
181
|
+
# generic kernel used to automatically generate overloads from launch args
|
|
182
|
+
@wp.kernel
|
|
183
|
+
def generic_fill(a: wp.array(dtype=Any), value: Any):
|
|
184
|
+
tid = wp.tid()
|
|
185
|
+
a[tid] = value
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
def test_generic_fill(test, device):
|
|
189
|
+
with wp.ScopedDevice(device):
|
|
190
|
+
n = 10
|
|
191
|
+
ai = wp.zeros(n, dtype=int)
|
|
192
|
+
af = wp.zeros(n, dtype=float)
|
|
193
|
+
a3 = wp.zeros(n, dtype=wp.vec3)
|
|
194
|
+
ab = wp.zeros(n, dtype=wp.bool)
|
|
195
|
+
|
|
196
|
+
wp.launch(generic_fill, dim=ai.size, inputs=[ai, 42])
|
|
197
|
+
wp.launch(generic_fill, dim=af.size, inputs=[af, 17.0])
|
|
198
|
+
wp.launch(generic_fill, dim=a3.size, inputs=[a3, wp.vec3(5.0, 5.0, 5.0)])
|
|
199
|
+
wp.launch(generic_fill, dim=ab.size, inputs=[ab, wp.bool(True)])
|
|
200
|
+
|
|
201
|
+
assert_np_equal(ai.numpy(), np.full((n,), 42, dtype=np.int32))
|
|
202
|
+
assert_np_equal(af.numpy(), np.full((n,), 17.0, dtype=np.float32))
|
|
203
|
+
assert_np_equal(a3.numpy(), np.full((n, 3), 5.0, dtype=np.float32))
|
|
204
|
+
assert_np_equal(ab.numpy(), np.full((n,), True, dtype=np.bool_))
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
# generic kernel used to create and launch explicit overloads
|
|
208
|
+
@wp.kernel
|
|
209
|
+
def generic_fill_v2(a: wp.array(dtype=Any), value: Any):
|
|
210
|
+
tid = wp.tid()
|
|
211
|
+
a[tid] = value
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
vec3b_type = wp.vec(3, wp.bool)
|
|
215
|
+
# create explicit overloads to be launched directly
|
|
216
|
+
fill_int = wp.overload(generic_fill_v2, [wp.array(dtype=int), int])
|
|
217
|
+
fill_float = wp.overload(generic_fill_v2, [wp.array(dtype=float), float])
|
|
218
|
+
fill_vec3 = wp.overload(generic_fill_v2, [wp.array(dtype=wp.vec3), wp.vec3])
|
|
219
|
+
fill_vec3b = wp.overload(generic_fill_v2, [wp.array(dtype=vec3b_type), vec3b_type])
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
def test_generic_fill_overloads(test, device):
|
|
223
|
+
with wp.ScopedDevice(device):
|
|
224
|
+
n = 10
|
|
225
|
+
ai = wp.zeros(n, dtype=int)
|
|
226
|
+
af = wp.zeros(n, dtype=float)
|
|
227
|
+
a3 = wp.zeros(n, dtype=wp.vec3)
|
|
228
|
+
a3b = wp.zeros(n, dtype=vec3b_type)
|
|
229
|
+
|
|
230
|
+
wp.launch(fill_int, dim=ai.size, inputs=[ai, 42])
|
|
231
|
+
wp.launch(fill_float, dim=af.size, inputs=[af, 17.0])
|
|
232
|
+
wp.launch(fill_vec3, dim=a3.size, inputs=[a3, wp.vec3(5.0, 5.0, 5.0)])
|
|
233
|
+
wp.launch(fill_vec3b, dim=a3b.size, inputs=[a3b, vec3b_type([True, True, True])])
|
|
234
|
+
|
|
235
|
+
assert_np_equal(ai.numpy(), np.full((n,), 42, dtype=np.int32))
|
|
236
|
+
assert_np_equal(af.numpy(), np.full((n,), 17.0, dtype=np.float32))
|
|
237
|
+
assert_np_equal(a3.numpy(), np.full((n, 3), 5.0, dtype=np.float32))
|
|
238
|
+
assert_np_equal(a3b.numpy(), np.full((n, 3), True, dtype=np.bool_))
|
|
239
|
+
|
|
240
|
+
|
|
241
|
+
# generic kernel used to test generic types mixed with specialized types
|
|
242
|
+
@wp.func
|
|
243
|
+
def generic_conditional_setter_func(a: wp.array(dtype=Any), i: int, value: Any, relative: bool):
|
|
244
|
+
if relative:
|
|
245
|
+
a[i] += value
|
|
246
|
+
else:
|
|
247
|
+
a[i] = value
|
|
248
|
+
|
|
249
|
+
|
|
250
|
+
@wp.kernel
|
|
251
|
+
def generic_conditional_setter(a: wp.array(dtype=Any), i: int, value: Any, relative: bool):
|
|
252
|
+
generic_conditional_setter_func(a, i, value, relative)
|
|
253
|
+
|
|
254
|
+
|
|
255
|
+
def test_generic_conditional_setter(test, device):
|
|
256
|
+
with wp.ScopedDevice(device):
|
|
257
|
+
n = 10
|
|
258
|
+
ai = wp.zeros(n, dtype=int)
|
|
259
|
+
|
|
260
|
+
wp.launch(generic_conditional_setter, dim=1, inputs=[ai, 1, 42, False])
|
|
261
|
+
wp.launch(generic_conditional_setter, dim=1, inputs=[ai, 1, 5, True])
|
|
262
|
+
ai_true = np.zeros(n, dtype=np.int32)
|
|
263
|
+
ai_true[1] = 47
|
|
264
|
+
assert_np_equal(ai.numpy(), ai_true)
|
|
265
|
+
|
|
266
|
+
|
|
267
|
+
# custom vector/matrix types
|
|
268
|
+
my_vec5 = wp.vec(length=5, dtype=wp.float32)
|
|
269
|
+
my_mat55 = wp.mat(shape=(5, 5), dtype=wp.float32)
|
|
270
|
+
|
|
271
|
+
|
|
272
|
+
@wp.kernel
|
|
273
|
+
def generic_transform(v: Any, m: Any, expected: Any):
|
|
274
|
+
result = wp.mul(m, v)
|
|
275
|
+
wp.expect_eq(result, expected)
|
|
276
|
+
|
|
277
|
+
|
|
278
|
+
# use overload decorator syntax
|
|
279
|
+
@wp.overload
|
|
280
|
+
def generic_transform(v: wp.vec2, m: wp.mat22, expected: wp.vec2): # fmt: skip
|
|
281
|
+
...
|
|
282
|
+
|
|
283
|
+
|
|
284
|
+
@wp.overload
|
|
285
|
+
def generic_transform(v: wp.vec3, m: wp.mat33, expected: wp.vec3): # fmt: skip
|
|
286
|
+
...
|
|
287
|
+
|
|
288
|
+
|
|
289
|
+
@wp.overload
|
|
290
|
+
def generic_transform(v: wp.vec4, m: wp.mat44, expected: wp.vec4): # fmt: skip
|
|
291
|
+
...
|
|
292
|
+
|
|
293
|
+
|
|
294
|
+
@wp.overload
|
|
295
|
+
def generic_transform(v: my_vec5, m: my_mat55, expected: my_vec5): # fmt: skip
|
|
296
|
+
...
|
|
297
|
+
|
|
298
|
+
|
|
299
|
+
def test_generic_transform_kernel(test, device):
|
|
300
|
+
with wp.ScopedDevice(device):
|
|
301
|
+
v2 = wp.vec2(1, 2)
|
|
302
|
+
m22 = wp.mat22(2, 0, 0, 2)
|
|
303
|
+
e2 = wp.vec2(2, 4)
|
|
304
|
+
|
|
305
|
+
v3 = wp.vec3(1, 2, 3)
|
|
306
|
+
m33 = wp.mat33(2, 0, 0, 0, 2, 0, 0, 0, 2)
|
|
307
|
+
e3 = wp.vec3(2, 4, 6)
|
|
308
|
+
|
|
309
|
+
v4 = wp.vec4(1, 2, 3, 4)
|
|
310
|
+
m44 = wp.mat44(2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 2)
|
|
311
|
+
e4 = wp.vec4(2, 4, 6, 8)
|
|
312
|
+
|
|
313
|
+
v5 = my_vec5(1, 2, 3, 4, 5)
|
|
314
|
+
m55 = my_mat55(2, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 2)
|
|
315
|
+
e5 = my_vec5(2, 4, 6, 8, 10)
|
|
316
|
+
|
|
317
|
+
wp.launch(generic_transform, dim=1, inputs=[v2, m22, e2])
|
|
318
|
+
wp.launch(generic_transform, dim=1, inputs=[v3, m33, e3])
|
|
319
|
+
wp.launch(generic_transform, dim=1, inputs=[v4, m44, e4])
|
|
320
|
+
wp.launch(generic_transform, dim=1, inputs=[v5, m55, e5])
|
|
321
|
+
|
|
322
|
+
wp.synchronize()
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
@wp.kernel
|
|
326
|
+
def generic_transform_array(v: wp.array(), m: wp.array(), result: wp.array()):
|
|
327
|
+
tid = wp.tid()
|
|
328
|
+
result[tid] = wp.mul(m[tid], v[tid])
|
|
329
|
+
|
|
330
|
+
|
|
331
|
+
wp.overload(generic_transform_array, [wp.array(dtype=wp.vec2), wp.array(dtype=wp.mat22), wp.array(dtype=wp.vec2)])
|
|
332
|
+
wp.overload(generic_transform_array, [wp.array(dtype=wp.vec3), wp.array(dtype=wp.mat33), wp.array(dtype=wp.vec3)])
|
|
333
|
+
wp.overload(generic_transform_array, [wp.array(dtype=wp.vec4), wp.array(dtype=wp.mat44), wp.array(dtype=wp.vec4)])
|
|
334
|
+
wp.overload(generic_transform_array, [wp.array(dtype=my_vec5), wp.array(dtype=my_mat55), wp.array(dtype=my_vec5)])
|
|
335
|
+
|
|
336
|
+
|
|
337
|
+
def test_generic_transform_array_kernel(test, device):
|
|
338
|
+
with wp.ScopedDevice(device):
|
|
339
|
+
n = 10
|
|
340
|
+
|
|
341
|
+
a2_data = np.tile(np.arange(2, dtype=np.float32), (n, 1))
|
|
342
|
+
a3_data = np.tile(np.arange(3, dtype=np.float32), (n, 1))
|
|
343
|
+
a4_data = np.tile(np.arange(4, dtype=np.float32), (n, 1))
|
|
344
|
+
a5_data = np.tile(np.arange(5, dtype=np.float32), (n, 1))
|
|
345
|
+
|
|
346
|
+
m22_data = np.tile((np.identity(2, dtype=np.float32) * 2), (n, 1, 1))
|
|
347
|
+
m33_data = np.tile((np.identity(3, dtype=np.float32) * 2), (n, 1, 1))
|
|
348
|
+
m44_data = np.tile((np.identity(4, dtype=np.float32) * 2), (n, 1, 1))
|
|
349
|
+
m55_data = np.tile((np.identity(5, dtype=np.float32) * 2), (n, 1, 1))
|
|
350
|
+
|
|
351
|
+
a2 = wp.array(data=a2_data, dtype=wp.vec2)
|
|
352
|
+
a3 = wp.array(data=a3_data, dtype=wp.vec3)
|
|
353
|
+
a4 = wp.array(data=a4_data, dtype=wp.vec4)
|
|
354
|
+
a5 = wp.array(data=a5_data, dtype=my_vec5)
|
|
355
|
+
|
|
356
|
+
m22 = wp.array(data=m22_data, dtype=wp.mat22)
|
|
357
|
+
m33 = wp.array(data=m33_data, dtype=wp.mat33)
|
|
358
|
+
m44 = wp.array(data=m44_data, dtype=wp.mat44)
|
|
359
|
+
m55 = wp.array(data=m55_data, dtype=my_mat55)
|
|
360
|
+
|
|
361
|
+
b2 = wp.zeros_like(a2)
|
|
362
|
+
b3 = wp.zeros_like(a3)
|
|
363
|
+
b4 = wp.zeros_like(a4)
|
|
364
|
+
b5 = wp.zeros_like(a5)
|
|
365
|
+
|
|
366
|
+
wp.launch(generic_transform_array, dim=n, inputs=[a2, m22, b2])
|
|
367
|
+
wp.launch(generic_transform_array, dim=n, inputs=[a3, m33, b3])
|
|
368
|
+
wp.launch(generic_transform_array, dim=n, inputs=[a4, m44, b4])
|
|
369
|
+
wp.launch(generic_transform_array, dim=n, inputs=[a5, m55, b5])
|
|
370
|
+
|
|
371
|
+
assert_np_equal(b2.numpy(), a2_data * 2)
|
|
372
|
+
assert_np_equal(b3.numpy(), a3_data * 2)
|
|
373
|
+
assert_np_equal(b4.numpy(), a4_data * 2)
|
|
374
|
+
assert_np_equal(b5.numpy(), a5_data * 2)
|
|
375
|
+
|
|
376
|
+
|
|
377
|
+
@wp.struct
|
|
378
|
+
class Foo:
|
|
379
|
+
x: float
|
|
380
|
+
y: float
|
|
381
|
+
z: float
|
|
382
|
+
|
|
383
|
+
|
|
384
|
+
@wp.struct
|
|
385
|
+
class Bar:
|
|
386
|
+
x: wp.vec3
|
|
387
|
+
y: wp.vec3
|
|
388
|
+
z: wp.vec3
|
|
389
|
+
|
|
390
|
+
|
|
391
|
+
@wp.kernel
|
|
392
|
+
def test_generic_struct_kernel(s: Any):
|
|
393
|
+
# test member access for generic structs
|
|
394
|
+
wp.expect_eq(s.x + s.y, s.z)
|
|
395
|
+
|
|
396
|
+
|
|
397
|
+
wp.overload(test_generic_struct_kernel, [Foo])
|
|
398
|
+
wp.overload(test_generic_struct_kernel, [Bar])
|
|
399
|
+
|
|
400
|
+
|
|
401
|
+
@wp.kernel
|
|
402
|
+
def test_generic_type_cast_kernel(a: Any, b: Any):
|
|
403
|
+
a = type(a)(b)
|
|
404
|
+
c = type(generic_adder(b, b))(a)
|
|
405
|
+
wp.expect_eq(b, c)
|
|
406
|
+
|
|
407
|
+
|
|
408
|
+
wp.overload(test_generic_type_cast_kernel, [wp.float32, wp.float64])
|
|
409
|
+
wp.overload(test_generic_type_cast_kernel, [wp.float32, wp.int32])
|
|
410
|
+
wp.overload(test_generic_type_cast_kernel, [wp.vec3f, wp.vec3d])
|
|
411
|
+
wp.overload(test_generic_type_cast_kernel, [wp.mat22f, wp.mat22d])
|
|
412
|
+
|
|
413
|
+
|
|
414
|
+
def test_generic_type_cast(test, device):
|
|
415
|
+
with wp.ScopedDevice(device):
|
|
416
|
+
wp.launch(test_generic_type_cast_kernel, dim=1, inputs=[1.0, 2.0])
|
|
417
|
+
wp.launch(test_generic_type_cast_kernel, dim=1, inputs=[2.0, -5])
|
|
418
|
+
wp.launch(test_generic_type_cast_kernel, dim=1, inputs=[wp.vec3f(1.0, 2.0, 3.0), wp.vec3d(4.0, 5.0, 6.0)])
|
|
419
|
+
wp.launch(test_generic_type_cast_kernel, dim=1, inputs=[wp.mat22f(0.0), wp.mat22d(np.eye(2))])
|
|
420
|
+
|
|
421
|
+
wp.synchronize()
|
|
422
|
+
|
|
423
|
+
|
|
424
|
+
@wp.kernel
|
|
425
|
+
def test_generic_scalar_construction_kernel(a: wp.array(dtype=Any)):
|
|
426
|
+
zero = type(a[0])(0)
|
|
427
|
+
copy = a.dtype(a[0])
|
|
428
|
+
copy += zero
|
|
429
|
+
wp.expect_eq(copy, a[0])
|
|
430
|
+
|
|
431
|
+
|
|
432
|
+
wp.overload(test_generic_scalar_construction_kernel, [wp.array(dtype=wp.int32)])
|
|
433
|
+
wp.overload(test_generic_scalar_construction_kernel, [wp.array(dtype=wp.float64)])
|
|
434
|
+
|
|
435
|
+
|
|
436
|
+
def test_generic_scalar_construction(test, device):
|
|
437
|
+
with wp.ScopedDevice(device):
|
|
438
|
+
wp.launch(test_generic_scalar_construction_kernel, dim=1, inputs=[wp.array([1.0], dtype=wp.int32)])
|
|
439
|
+
wp.launch(test_generic_scalar_construction_kernel, dim=1, inputs=[wp.array([-5], dtype=wp.float64)])
|
|
440
|
+
|
|
441
|
+
wp.synchronize()
|
|
442
|
+
|
|
443
|
+
|
|
444
|
+
@wp.kernel
|
|
445
|
+
def test_generic_type_construction_kernel(a: wp.array(dtype=Any)):
|
|
446
|
+
zero = type(a[0])()
|
|
447
|
+
copy = type(a).dtype(a[0]) * a.dtype.dtype(1.0)
|
|
448
|
+
copy += zero
|
|
449
|
+
wp.expect_eq(copy, a[0])
|
|
450
|
+
|
|
451
|
+
|
|
452
|
+
wp.overload(test_generic_type_construction_kernel, [wp.array(dtype=wp.vec3f)])
|
|
453
|
+
wp.overload(test_generic_type_construction_kernel, [wp.array(dtype=wp.mat22d)])
|
|
454
|
+
|
|
455
|
+
|
|
456
|
+
def test_generic_type_construction(test, device):
|
|
457
|
+
with wp.ScopedDevice(device):
|
|
458
|
+
wp.launch(test_generic_type_construction_kernel, dim=1, inputs=[wp.array([1.0, 2.0, 3.0], dtype=wp.vec3f)])
|
|
459
|
+
wp.launch(test_generic_type_construction_kernel, dim=1, inputs=[wp.array([np.eye(2)], dtype=wp.mat22d)])
|
|
460
|
+
|
|
461
|
+
wp.synchronize()
|
|
462
|
+
|
|
463
|
+
|
|
464
|
+
@wp.kernel
|
|
465
|
+
def test_generic_struct_construction_kernel(a: Any):
|
|
466
|
+
b = type(a)(a.x, a.y, a.z)
|
|
467
|
+
wp.expect_eq(a.x, b.x)
|
|
468
|
+
wp.expect_eq(a.y, b.y)
|
|
469
|
+
wp.expect_eq(a.z, b.z)
|
|
470
|
+
|
|
471
|
+
|
|
472
|
+
wp.overload(test_generic_struct_construction_kernel, [Foo])
|
|
473
|
+
wp.overload(test_generic_struct_construction_kernel, [Bar])
|
|
474
|
+
|
|
475
|
+
|
|
476
|
+
@wp.kernel
|
|
477
|
+
def test_generic_type_as_argument_kernel(a: Any):
|
|
478
|
+
vec = wp.vector(length=2, dtype=type(a))
|
|
479
|
+
matrix = wp.identity(n=vec.length, dtype=vec.dtype) * a
|
|
480
|
+
wp.expect_eq(wp.trace(matrix), type(a)(2.0) * a)
|
|
481
|
+
|
|
482
|
+
|
|
483
|
+
wp.overload(test_generic_type_as_argument_kernel, [wp.float32])
|
|
484
|
+
wp.overload(test_generic_type_as_argument_kernel, [wp.float64])
|
|
485
|
+
|
|
486
|
+
|
|
487
|
+
def test_generic_type_as_argument(test, device):
|
|
488
|
+
with wp.ScopedDevice(device):
|
|
489
|
+
wp.launch(test_generic_type_as_argument_kernel, dim=1, inputs=[2.0])
|
|
490
|
+
wp.launch(test_generic_type_as_argument_kernel, dim=1, inputs=[-1.0])
|
|
491
|
+
|
|
492
|
+
wp.synchronize()
|
|
493
|
+
|
|
494
|
+
|
|
495
|
+
def test_type_operator_misspell(test, device):
|
|
496
|
+
@wp.kernel
|
|
497
|
+
def kernel():
|
|
498
|
+
i = wp.tid()
|
|
499
|
+
_ = typez(i)(0)
|
|
500
|
+
|
|
501
|
+
with test.assertRaisesRegex(KeyError, r"Referencing undefined symbol: typez"):
|
|
502
|
+
wp.launch(
|
|
503
|
+
kernel,
|
|
504
|
+
dim=1,
|
|
505
|
+
inputs=[],
|
|
506
|
+
device=device,
|
|
507
|
+
)
|
|
508
|
+
|
|
509
|
+
|
|
510
|
+
def test_type_attribute_error(test, device):
|
|
511
|
+
@wp.kernel
|
|
512
|
+
def kernel():
|
|
513
|
+
a = wp.vec3(0.0)
|
|
514
|
+
_ = a.dtype.shape
|
|
515
|
+
|
|
516
|
+
with test.assertRaisesRegex(AttributeError, r"`shape` is not an attribute of '<class 'warp.types.float32'>'"):
|
|
517
|
+
wp.launch(
|
|
518
|
+
kernel,
|
|
519
|
+
dim=1,
|
|
520
|
+
inputs=[],
|
|
521
|
+
device=device,
|
|
522
|
+
)
|
|
523
|
+
|
|
524
|
+
|
|
525
|
+
class TestGenerics(unittest.TestCase):
|
|
526
|
+
pass
|
|
527
|
+
|
|
528
|
+
|
|
529
|
+
devices = get_test_devices()
|
|
530
|
+
|
|
531
|
+
add_kernel_test(TestGenerics, name="test_generic_adder", kernel=test_generic_adder, dim=1, devices=devices)
|
|
532
|
+
add_kernel_test(TestGenerics, name="test_specialized_func", kernel=test_specialized_func, dim=1, devices=devices)
|
|
533
|
+
|
|
534
|
+
add_function_test(TestGenerics, "test_generic_array_kernel", test_generic_array_kernel, devices=devices)
|
|
535
|
+
add_function_test(TestGenerics, "test_generic_accumulator_kernel", test_generic_accumulator_kernel, devices=devices)
|
|
536
|
+
add_function_test(TestGenerics, "test_generic_fill", test_generic_fill, devices=devices)
|
|
537
|
+
add_function_test(TestGenerics, "test_generic_fill_overloads", test_generic_fill_overloads, devices=devices)
|
|
538
|
+
add_function_test(TestGenerics, "test_generic_conditional_setter", test_generic_conditional_setter, devices=devices)
|
|
539
|
+
add_function_test(TestGenerics, "test_generic_transform_kernel", test_generic_transform_kernel, devices=devices)
|
|
540
|
+
add_function_test(
|
|
541
|
+
TestGenerics, "test_generic_transform_array_kernel", test_generic_transform_array_kernel, devices=devices
|
|
542
|
+
)
|
|
543
|
+
add_function_test(TestGenerics, "test_generic_type_cast", test_generic_type_cast, devices=devices)
|
|
544
|
+
add_function_test(TestGenerics, "test_generic_type_construction", test_generic_type_construction, devices=devices)
|
|
545
|
+
add_function_test(TestGenerics, "test_generic_scalar_construction", test_generic_scalar_construction, devices=devices)
|
|
546
|
+
add_function_test(TestGenerics, "test_generic_type_as_argument", test_generic_type_as_argument, devices=devices)
|
|
547
|
+
|
|
548
|
+
foo = Foo()
|
|
549
|
+
foo.x = 17.0
|
|
550
|
+
foo.y = 25.0
|
|
551
|
+
foo.z = 42.0
|
|
552
|
+
|
|
553
|
+
bar = Bar()
|
|
554
|
+
bar.x = wp.vec3(1, 2, 3)
|
|
555
|
+
bar.y = wp.vec3(10, 20, 30)
|
|
556
|
+
bar.z = wp.vec3(11, 22, 33)
|
|
557
|
+
|
|
558
|
+
add_kernel_test(
|
|
559
|
+
TestGenerics,
|
|
560
|
+
name="test_generic_struct_kernel",
|
|
561
|
+
kernel=test_generic_struct_kernel,
|
|
562
|
+
dim=1,
|
|
563
|
+
inputs=[foo],
|
|
564
|
+
devices=devices,
|
|
565
|
+
)
|
|
566
|
+
add_kernel_test(
|
|
567
|
+
TestGenerics,
|
|
568
|
+
name="test_generic_struct_kernel",
|
|
569
|
+
kernel=test_generic_struct_kernel,
|
|
570
|
+
dim=1,
|
|
571
|
+
inputs=[bar],
|
|
572
|
+
devices=devices,
|
|
573
|
+
)
|
|
574
|
+
|
|
575
|
+
add_kernel_test(
|
|
576
|
+
TestGenerics,
|
|
577
|
+
name="test_generic_struct_construction_kernel",
|
|
578
|
+
kernel=test_generic_struct_construction_kernel,
|
|
579
|
+
dim=1,
|
|
580
|
+
inputs=[foo],
|
|
581
|
+
devices=devices,
|
|
582
|
+
)
|
|
583
|
+
add_kernel_test(
|
|
584
|
+
TestGenerics,
|
|
585
|
+
name="test_generic_struct_construction_kernel",
|
|
586
|
+
kernel=test_generic_struct_construction_kernel,
|
|
587
|
+
dim=1,
|
|
588
|
+
inputs=[bar],
|
|
589
|
+
devices=devices,
|
|
590
|
+
)
|
|
591
|
+
add_function_test(TestGenerics, "test_type_operator_misspell", test_type_operator_misspell, devices=devices)
|
|
592
|
+
add_function_test(TestGenerics, "test_type_attribute_error", test_type_attribute_error, devices=devices)
|
|
593
|
+
|
|
594
|
+
if __name__ == "__main__":
|
|
595
|
+
wp.build.clear_kernel_cache()
|
|
596
|
+
unittest.main(verbosity=2)
|