warp-lang 1.9.0__py3-none-win_amd64.whl → 1.10.0rc2__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 +301 -287
- warp/__init__.pyi +2220 -313
- warp/_src/__init__.py +14 -0
- warp/_src/autograd.py +1075 -0
- warp/_src/build.py +618 -0
- warp/_src/build_dll.py +640 -0
- warp/{builtins.py → _src/builtins.py} +1497 -226
- warp/_src/codegen.py +4359 -0
- warp/{config.py → _src/config.py} +178 -169
- warp/_src/constants.py +57 -0
- warp/_src/context.py +8294 -0
- warp/_src/dlpack.py +462 -0
- warp/_src/fabric.py +355 -0
- warp/_src/fem/__init__.py +14 -0
- warp/_src/fem/adaptivity.py +508 -0
- warp/_src/fem/cache.py +687 -0
- warp/_src/fem/dirichlet.py +188 -0
- warp/{fem → _src/fem}/domain.py +40 -30
- warp/_src/fem/field/__init__.py +131 -0
- warp/_src/fem/field/field.py +701 -0
- warp/{fem → _src/fem}/field/nodal_field.py +30 -15
- warp/{fem → _src/fem}/field/restriction.py +1 -1
- warp/{fem → _src/fem}/field/virtual.py +53 -27
- warp/_src/fem/geometry/__init__.py +32 -0
- warp/{fem → _src/fem}/geometry/adaptive_nanogrid.py +77 -163
- warp/_src/fem/geometry/closest_point.py +97 -0
- warp/{fem → _src/fem}/geometry/deformed_geometry.py +14 -22
- warp/{fem → _src/fem}/geometry/element.py +32 -10
- warp/{fem → _src/fem}/geometry/geometry.py +48 -20
- warp/{fem → _src/fem}/geometry/grid_2d.py +12 -23
- warp/{fem → _src/fem}/geometry/grid_3d.py +12 -23
- warp/{fem → _src/fem}/geometry/hexmesh.py +40 -63
- warp/{fem → _src/fem}/geometry/nanogrid.py +255 -248
- warp/{fem → _src/fem}/geometry/partition.py +121 -63
- warp/{fem → _src/fem}/geometry/quadmesh.py +26 -45
- warp/{fem → _src/fem}/geometry/tetmesh.py +40 -63
- warp/{fem → _src/fem}/geometry/trimesh.py +26 -45
- warp/{fem → _src/fem}/integrate.py +164 -158
- warp/_src/fem/linalg.py +383 -0
- warp/_src/fem/operator.py +396 -0
- warp/_src/fem/polynomial.py +229 -0
- warp/{fem → _src/fem}/quadrature/pic_quadrature.py +15 -20
- warp/{fem → _src/fem}/quadrature/quadrature.py +95 -47
- warp/_src/fem/space/__init__.py +248 -0
- warp/{fem → _src/fem}/space/basis_function_space.py +20 -11
- warp/_src/fem/space/basis_space.py +679 -0
- warp/{fem → _src/fem}/space/dof_mapper.py +3 -3
- warp/{fem → _src/fem}/space/function_space.py +14 -13
- warp/{fem → _src/fem}/space/grid_2d_function_space.py +4 -7
- warp/{fem → _src/fem}/space/grid_3d_function_space.py +4 -4
- warp/{fem → _src/fem}/space/hexmesh_function_space.py +4 -10
- warp/{fem → _src/fem}/space/nanogrid_function_space.py +3 -9
- warp/{fem → _src/fem}/space/partition.py +117 -60
- warp/{fem → _src/fem}/space/quadmesh_function_space.py +4 -10
- warp/{fem → _src/fem}/space/restriction.py +66 -33
- warp/_src/fem/space/shape/__init__.py +152 -0
- warp/{fem → _src/fem}/space/shape/cube_shape_function.py +9 -9
- warp/{fem → _src/fem}/space/shape/shape_function.py +8 -9
- warp/{fem → _src/fem}/space/shape/square_shape_function.py +6 -6
- warp/{fem → _src/fem}/space/shape/tet_shape_function.py +3 -3
- warp/{fem → _src/fem}/space/shape/triangle_shape_function.py +3 -3
- warp/{fem → _src/fem}/space/tetmesh_function_space.py +3 -9
- warp/_src/fem/space/topology.py +459 -0
- warp/{fem → _src/fem}/space/trimesh_function_space.py +3 -9
- warp/_src/fem/types.py +112 -0
- warp/_src/fem/utils.py +486 -0
- warp/_src/jax.py +186 -0
- warp/_src/jax_experimental/__init__.py +14 -0
- warp/_src/jax_experimental/custom_call.py +387 -0
- warp/_src/jax_experimental/ffi.py +1284 -0
- warp/_src/jax_experimental/xla_ffi.py +656 -0
- warp/_src/marching_cubes.py +708 -0
- warp/_src/math.py +414 -0
- warp/_src/optim/__init__.py +14 -0
- warp/_src/optim/adam.py +163 -0
- warp/_src/optim/linear.py +1606 -0
- warp/_src/optim/sgd.py +112 -0
- warp/_src/paddle.py +406 -0
- warp/_src/render/__init__.py +14 -0
- warp/_src/render/imgui_manager.py +289 -0
- warp/_src/render/render_opengl.py +3636 -0
- warp/_src/render/render_usd.py +937 -0
- warp/_src/render/utils.py +160 -0
- warp/_src/sparse.py +2716 -0
- warp/_src/tape.py +1206 -0
- warp/{thirdparty → _src/thirdparty}/unittest_parallel.py +9 -2
- warp/_src/torch.py +391 -0
- warp/_src/types.py +5870 -0
- warp/_src/utils.py +1693 -0
- warp/autograd.py +12 -1054
- warp/bin/warp-clang.dll +0 -0
- warp/bin/warp.dll +0 -0
- warp/build.py +8 -588
- warp/build_dll.py +6 -471
- warp/codegen.py +6 -4246
- warp/constants.py +6 -39
- warp/context.py +12 -7851
- warp/dlpack.py +6 -444
- warp/examples/distributed/example_jacobi_mpi.py +4 -5
- warp/examples/fem/example_adaptive_grid.py +1 -1
- warp/examples/fem/example_apic_fluid.py +1 -1
- warp/examples/fem/example_burgers.py +8 -8
- warp/examples/fem/example_diffusion.py +1 -1
- warp/examples/fem/example_distortion_energy.py +1 -1
- warp/examples/fem/example_mixed_elasticity.py +2 -2
- warp/examples/fem/example_navier_stokes.py +1 -1
- warp/examples/fem/example_nonconforming_contact.py +7 -7
- warp/examples/fem/example_stokes.py +1 -1
- warp/examples/fem/example_stokes_transfer.py +1 -1
- warp/examples/fem/utils.py +2 -2
- warp/examples/interop/example_jax_callable.py +1 -1
- warp/examples/interop/example_jax_ffi_callback.py +1 -1
- warp/examples/interop/example_jax_kernel.py +3 -2
- warp/examples/tile/example_tile_mcgp.py +191 -0
- warp/fabric.py +6 -337
- warp/fem/__init__.py +159 -97
- warp/fem/adaptivity.py +7 -489
- warp/fem/cache.py +9 -648
- warp/fem/dirichlet.py +6 -184
- warp/fem/field/__init__.py +8 -109
- warp/fem/field/field.py +7 -652
- warp/fem/geometry/__init__.py +7 -18
- warp/fem/geometry/closest_point.py +11 -77
- warp/fem/linalg.py +18 -366
- warp/fem/operator.py +11 -369
- warp/fem/polynomial.py +9 -209
- warp/fem/space/__init__.py +5 -211
- warp/fem/space/basis_space.py +6 -662
- warp/fem/space/shape/__init__.py +41 -118
- warp/fem/space/topology.py +6 -437
- warp/fem/types.py +6 -81
- warp/fem/utils.py +11 -444
- warp/jax.py +8 -165
- warp/jax_experimental/__init__.py +14 -1
- warp/jax_experimental/custom_call.py +8 -342
- warp/jax_experimental/ffi.py +17 -853
- warp/jax_experimental/xla_ffi.py +5 -596
- warp/marching_cubes.py +5 -689
- warp/math.py +16 -393
- warp/native/array.h +385 -37
- warp/native/builtin.h +316 -39
- warp/native/bvh.cpp +43 -9
- warp/native/bvh.cu +62 -27
- warp/native/bvh.h +310 -309
- warp/native/clang/clang.cpp +102 -97
- warp/native/coloring.cpp +0 -1
- warp/native/crt.h +208 -0
- warp/native/exports.h +156 -0
- warp/native/hashgrid.cu +2 -0
- warp/native/intersect.h +24 -1
- warp/native/intersect_tri.h +44 -35
- warp/native/mat.h +1456 -276
- warp/native/mesh.cpp +4 -4
- warp/native/mesh.cu +4 -2
- warp/native/mesh.h +176 -61
- warp/native/quat.h +0 -52
- warp/native/scan.cu +2 -0
- warp/native/sort.cu +22 -13
- warp/native/sort.h +2 -0
- warp/native/sparse.cu +7 -3
- warp/native/spatial.h +12 -0
- warp/native/tile.h +837 -70
- warp/native/tile_radix_sort.h +1 -1
- warp/native/tile_reduce.h +394 -46
- warp/native/tile_scan.h +4 -4
- warp/native/vec.h +469 -53
- warp/native/version.h +23 -0
- warp/native/volume.cpp +1 -1
- warp/native/volume.cu +1 -0
- warp/native/volume.h +1 -1
- warp/native/volume_builder.cu +2 -0
- warp/native/warp.cpp +60 -32
- warp/native/warp.cu +313 -201
- warp/native/warp.h +14 -11
- warp/optim/__init__.py +6 -3
- warp/optim/adam.py +6 -145
- warp/optim/linear.py +14 -1585
- warp/optim/sgd.py +6 -94
- warp/paddle.py +6 -388
- warp/render/__init__.py +8 -4
- warp/render/imgui_manager.py +7 -267
- warp/render/render_opengl.py +6 -3616
- warp/render/render_usd.py +6 -918
- warp/render/utils.py +6 -142
- warp/sparse.py +37 -2563
- warp/tape.py +6 -1188
- warp/tests/__main__.py +1 -1
- warp/tests/cuda/test_async.py +4 -4
- warp/tests/cuda/test_conditional_captures.py +1 -1
- warp/tests/cuda/test_multigpu.py +1 -1
- warp/tests/cuda/test_streams.py +58 -1
- warp/tests/geometry/test_bvh.py +157 -22
- warp/tests/geometry/test_hash_grid.py +38 -0
- warp/tests/geometry/test_marching_cubes.py +0 -1
- warp/tests/geometry/test_mesh.py +5 -3
- warp/tests/geometry/test_mesh_query_aabb.py +5 -12
- warp/tests/geometry/test_mesh_query_point.py +5 -2
- warp/tests/geometry/test_mesh_query_ray.py +15 -3
- warp/tests/geometry/test_volume_write.py +5 -5
- warp/tests/interop/test_dlpack.py +14 -14
- warp/tests/interop/test_jax.py +1382 -79
- warp/tests/interop/test_paddle.py +1 -1
- warp/tests/test_adam.py +0 -1
- warp/tests/test_arithmetic.py +9 -9
- warp/tests/test_array.py +529 -100
- warp/tests/test_array_reduce.py +3 -3
- warp/tests/test_atomic.py +12 -8
- warp/tests/test_atomic_bitwise.py +209 -0
- warp/tests/test_atomic_cas.py +4 -4
- warp/tests/test_bool.py +2 -2
- warp/tests/test_builtins_resolution.py +5 -571
- warp/tests/test_codegen.py +34 -15
- warp/tests/test_conditional.py +1 -1
- warp/tests/test_context.py +6 -6
- warp/tests/test_copy.py +242 -161
- warp/tests/test_ctypes.py +3 -3
- warp/tests/test_devices.py +24 -2
- warp/tests/test_examples.py +16 -84
- warp/tests/test_fabricarray.py +35 -35
- warp/tests/test_fast_math.py +0 -2
- warp/tests/test_fem.py +60 -14
- warp/tests/test_fixedarray.py +3 -3
- warp/tests/test_func.py +8 -5
- warp/tests/test_generics.py +1 -1
- warp/tests/test_indexedarray.py +24 -24
- warp/tests/test_intersect.py +39 -9
- warp/tests/test_large.py +1 -1
- warp/tests/test_lerp.py +3 -1
- warp/tests/test_linear_solvers.py +1 -1
- warp/tests/test_map.py +49 -4
- warp/tests/test_mat.py +52 -62
- warp/tests/test_mat_constructors.py +4 -5
- warp/tests/test_mat_lite.py +1 -1
- warp/tests/test_mat_scalar_ops.py +121 -121
- warp/tests/test_math.py +34 -0
- warp/tests/test_module_aot.py +4 -4
- warp/tests/test_modules_lite.py +28 -2
- warp/tests/test_print.py +11 -11
- warp/tests/test_quat.py +93 -58
- warp/tests/test_runlength_encode.py +1 -1
- warp/tests/test_scalar_ops.py +38 -10
- warp/tests/test_smoothstep.py +1 -1
- warp/tests/test_sparse.py +126 -15
- warp/tests/test_spatial.py +105 -87
- warp/tests/test_special_values.py +6 -6
- warp/tests/test_static.py +7 -7
- warp/tests/test_struct.py +13 -2
- warp/tests/test_triangle_closest_point.py +48 -1
- warp/tests/test_tuple.py +96 -0
- warp/tests/test_types.py +82 -9
- warp/tests/test_utils.py +52 -52
- warp/tests/test_vec.py +29 -29
- warp/tests/test_vec_constructors.py +5 -5
- warp/tests/test_vec_scalar_ops.py +97 -97
- warp/tests/test_version.py +75 -0
- warp/tests/tile/test_tile.py +239 -0
- warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
- warp/tests/tile/test_tile_cholesky.py +7 -4
- warp/tests/tile/test_tile_load.py +26 -2
- warp/tests/tile/test_tile_mathdx.py +3 -3
- warp/tests/tile/test_tile_matmul.py +1 -1
- warp/tests/tile/test_tile_mlp.py +2 -4
- warp/tests/tile/test_tile_reduce.py +214 -13
- warp/tests/unittest_suites.py +6 -14
- warp/tests/unittest_utils.py +10 -9
- warp/tests/walkthrough_debug.py +3 -1
- warp/torch.py +6 -373
- warp/types.py +29 -5750
- warp/utils.py +10 -1659
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/METADATA +47 -103
- warp_lang-1.10.0rc2.dist-info/RECORD +468 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
- warp_lang-1.10.0rc2.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
- warp/examples/assets/cartpole.urdf +0 -110
- warp/examples/assets/crazyflie.usd +0 -0
- warp/examples/assets/nv_ant.xml +0 -92
- warp/examples/assets/nv_humanoid.xml +0 -183
- warp/examples/assets/quadruped.urdf +0 -268
- warp/examples/optim/example_bounce.py +0 -266
- warp/examples/optim/example_cloth_throw.py +0 -228
- warp/examples/optim/example_drone.py +0 -870
- warp/examples/optim/example_inverse_kinematics.py +0 -182
- warp/examples/optim/example_inverse_kinematics_torch.py +0 -191
- warp/examples/optim/example_softbody_properties.py +0 -400
- warp/examples/optim/example_spring_cage.py +0 -245
- warp/examples/optim/example_trajectory.py +0 -227
- warp/examples/sim/example_cartpole.py +0 -143
- warp/examples/sim/example_cloth.py +0 -225
- warp/examples/sim/example_cloth_self_contact.py +0 -316
- warp/examples/sim/example_granular.py +0 -130
- warp/examples/sim/example_granular_collision_sdf.py +0 -202
- warp/examples/sim/example_jacobian_ik.py +0 -244
- warp/examples/sim/example_particle_chain.py +0 -124
- warp/examples/sim/example_quadruped.py +0 -203
- warp/examples/sim/example_rigid_chain.py +0 -203
- warp/examples/sim/example_rigid_contact.py +0 -195
- warp/examples/sim/example_rigid_force.py +0 -133
- warp/examples/sim/example_rigid_gyroscopic.py +0 -115
- warp/examples/sim/example_rigid_soft_contact.py +0 -140
- warp/examples/sim/example_soft_body.py +0 -196
- warp/examples/tile/example_tile_walker.py +0 -327
- warp/sim/__init__.py +0 -74
- warp/sim/articulation.py +0 -793
- warp/sim/collide.py +0 -2570
- warp/sim/graph_coloring.py +0 -307
- warp/sim/import_mjcf.py +0 -791
- warp/sim/import_snu.py +0 -227
- warp/sim/import_urdf.py +0 -579
- warp/sim/import_usd.py +0 -898
- warp/sim/inertia.py +0 -357
- warp/sim/integrator.py +0 -245
- warp/sim/integrator_euler.py +0 -2000
- warp/sim/integrator_featherstone.py +0 -2101
- warp/sim/integrator_vbd.py +0 -2487
- warp/sim/integrator_xpbd.py +0 -3295
- warp/sim/model.py +0 -4821
- warp/sim/particles.py +0 -121
- warp/sim/render.py +0 -431
- warp/sim/utils.py +0 -431
- warp/tests/sim/disabled_kinematics.py +0 -244
- warp/tests/sim/test_cloth.py +0 -863
- warp/tests/sim/test_collision.py +0 -743
- warp/tests/sim/test_coloring.py +0 -347
- warp/tests/sim/test_inertia.py +0 -161
- warp/tests/sim/test_model.py +0 -226
- warp/tests/sim/test_sim_grad.py +0 -287
- warp/tests/sim/test_sim_grad_bounce_linear.py +0 -212
- warp/tests/sim/test_sim_kinematics.py +0 -98
- warp/thirdparty/__init__.py +0 -0
- warp_lang-1.9.0.dist-info/RECORD +0 -456
- /warp/{fem → _src/fem}/quadrature/__init__.py +0 -0
- /warp/{tests/sim → _src/thirdparty}/__init__.py +0 -0
- /warp/{thirdparty → _src/thirdparty}/appdirs.py +0 -0
- /warp/{thirdparty → _src/thirdparty}/dlpack.py +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/WHEEL +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/licenses/LICENSE.md +0 -0
- {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/top_level.txt +0 -0
warp/native/warp.h
CHANGED
|
@@ -28,7 +28,7 @@ struct timing_result_t;
|
|
|
28
28
|
// this is the core runtime API exposed on the DLL level
|
|
29
29
|
extern "C"
|
|
30
30
|
{
|
|
31
|
-
WP_API int wp_init();
|
|
31
|
+
WP_API int wp_init(const char* expected_version);
|
|
32
32
|
//WP_API void wp_shutdown();
|
|
33
33
|
|
|
34
34
|
// get error message from C++
|
|
@@ -75,23 +75,25 @@ extern "C"
|
|
|
75
75
|
WP_API void wp_memtile_host(void* dest, const void* src, size_t srcsize, size_t n);
|
|
76
76
|
WP_API void wp_memtile_device(void* context, void* dest, const void* src, size_t srcsize, size_t n);
|
|
77
77
|
|
|
78
|
-
WP_API uint64_t wp_bvh_create_host(wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type);
|
|
78
|
+
WP_API uint64_t wp_bvh_create_host(wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type, int leaf_size);
|
|
79
79
|
WP_API void wp_bvh_destroy_host(uint64_t id);
|
|
80
80
|
WP_API void wp_bvh_refit_host(uint64_t id);
|
|
81
|
+
WP_API void wp_bvh_rebuild_host(uint64_t id, int constructor_type);
|
|
81
82
|
|
|
82
|
-
|
|
83
|
+
WP_API uint64_t wp_bvh_create_device(void* context, wp::vec3* lowers, wp::vec3* uppers, int num_items, int constructor_type, int leaf_size);
|
|
83
84
|
WP_API void wp_bvh_destroy_device(uint64_t id);
|
|
84
85
|
WP_API void wp_bvh_refit_device(uint64_t id);
|
|
86
|
+
WP_API void wp_bvh_rebuild_device(uint64_t id);
|
|
85
87
|
|
|
86
88
|
// create a user-accessible copy of the mesh, it is the
|
|
87
89
|
// users responsibility to keep-alive the points/tris data for the duration of the mesh lifetime
|
|
88
90
|
WP_API uint64_t wp_mesh_create_host(wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities, wp::array_t<int> tris,
|
|
89
|
-
int num_points, int num_tris, int support_winding_number, int constructor_type);
|
|
91
|
+
int num_points, int num_tris, int support_winding_number, int constructor_type, int bvh_leaf_size);
|
|
90
92
|
WP_API void wp_mesh_destroy_host(uint64_t id);
|
|
91
93
|
WP_API void wp_mesh_refit_host(uint64_t id);
|
|
92
94
|
|
|
93
95
|
WP_API uint64_t wp_mesh_create_device(void* context, wp::array_t<wp::vec3> points, wp::array_t<wp::vec3> velocities,
|
|
94
|
-
wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number, int constructor_type);
|
|
96
|
+
wp::array_t<int> tris, int num_points, int num_tris, int support_winding_number, int constructor_type, int bvh_leaf_size);
|
|
95
97
|
WP_API void wp_mesh_destroy_device(uint64_t id);
|
|
96
98
|
WP_API void wp_mesh_refit_device(uint64_t id);
|
|
97
99
|
|
|
@@ -232,6 +234,7 @@ extern "C"
|
|
|
232
234
|
|
|
233
235
|
WP_API int wp_cuda_driver_version(); // CUDA driver version
|
|
234
236
|
WP_API int wp_cuda_toolkit_version(); // CUDA Toolkit version used to build Warp
|
|
237
|
+
WP_API const char* wp_version(); // Warp native library version string
|
|
235
238
|
WP_API bool wp_cuda_driver_is_initialized();
|
|
236
239
|
|
|
237
240
|
WP_API int wp_nvrtc_supported_arch_count();
|
|
@@ -293,8 +296,8 @@ extern "C"
|
|
|
293
296
|
WP_API void wp_cuda_stream_unregister(void* context, void* stream);
|
|
294
297
|
WP_API void* wp_cuda_stream_get_current();
|
|
295
298
|
WP_API void wp_cuda_stream_synchronize(void* stream);
|
|
296
|
-
WP_API void wp_cuda_stream_wait_event(void* stream, void* event);
|
|
297
|
-
WP_API void wp_cuda_stream_wait_stream(void* stream, void* other_stream, void* event);
|
|
299
|
+
WP_API void wp_cuda_stream_wait_event(void* stream, void* event, bool external=false);
|
|
300
|
+
WP_API void wp_cuda_stream_wait_stream(void* stream, void* other_stream, void* event, bool external=false);
|
|
298
301
|
WP_API int wp_cuda_stream_is_capturing(void* stream);
|
|
299
302
|
WP_API uint64_t wp_cuda_stream_get_capture_id(void* stream);
|
|
300
303
|
WP_API int wp_cuda_stream_get_priority(void* stream);
|
|
@@ -302,7 +305,7 @@ extern "C"
|
|
|
302
305
|
WP_API void* wp_cuda_event_create(void* context, unsigned flags);
|
|
303
306
|
WP_API void wp_cuda_event_destroy(void* event);
|
|
304
307
|
WP_API int wp_cuda_event_query(void* event);
|
|
305
|
-
WP_API void wp_cuda_event_record(void* event, void* stream, bool
|
|
308
|
+
WP_API void wp_cuda_event_record(void* event, void* stream, bool external=false);
|
|
306
309
|
WP_API void wp_cuda_event_synchronize(void* event);
|
|
307
310
|
WP_API float wp_cuda_event_elapsed_time(void* start_event, void* end_event);
|
|
308
311
|
|
|
@@ -314,9 +317,9 @@ extern "C"
|
|
|
314
317
|
WP_API bool wp_cuda_graph_exec_destroy(void* context, void* graph_exec);
|
|
315
318
|
WP_API bool wp_capture_debug_dot_print(void* graph, const char *path, uint32_t flags);
|
|
316
319
|
|
|
317
|
-
WP_API bool wp_cuda_graph_insert_if_else(void* context, void* stream, int* condition, void** if_graph_ret, void** else_graph_ret);
|
|
318
|
-
WP_API bool wp_cuda_graph_insert_while(void* context, void* stream, int* condition, void** body_graph_ret, uint64_t* handle_ret);
|
|
319
|
-
WP_API bool wp_cuda_graph_set_condition(void* context, void* stream, int* condition, uint64_t handle);
|
|
320
|
+
WP_API bool wp_cuda_graph_insert_if_else(void* context, void* stream, int arch, bool use_ptx, int* condition, void** if_graph_ret, void** else_graph_ret);
|
|
321
|
+
WP_API bool wp_cuda_graph_insert_while(void* context, void* stream, int arch, bool use_ptx, int* condition, void** body_graph_ret, uint64_t* handle_ret);
|
|
322
|
+
WP_API bool wp_cuda_graph_set_condition(void* context, void* stream, int arch, bool use_ptx, int* condition, uint64_t handle);
|
|
320
323
|
WP_API bool wp_cuda_graph_pause_capture(void* context, void* stream, void** graph_ret);
|
|
321
324
|
WP_API bool wp_cuda_graph_resume_capture(void* context, void* stream, void* graph);
|
|
322
325
|
WP_API bool wp_cuda_graph_insert_child_graph(void* context, void* stream, void* child_graph);
|
warp/optim/__init__.py
CHANGED
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c)
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
2
|
# SPDX-License-Identifier: Apache-2.0
|
|
3
3
|
#
|
|
4
4
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
@@ -13,5 +13,8 @@
|
|
|
13
13
|
# See the License for the specific language governing permissions and
|
|
14
14
|
# limitations under the License.
|
|
15
15
|
|
|
16
|
-
|
|
17
|
-
|
|
16
|
+
# isort: skip_file
|
|
17
|
+
|
|
18
|
+
from warp._src.optim.adam import Adam as Adam
|
|
19
|
+
|
|
20
|
+
from warp._src.optim.sgd import SGD as SGD
|
warp/optim/adam.py
CHANGED
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
# SPDX-FileCopyrightText: Copyright (c)
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
2
2
|
# SPDX-License-Identifier: Apache-2.0
|
|
3
3
|
#
|
|
4
4
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
@@ -13,151 +13,12 @@
|
|
|
13
13
|
# See the License for the specific language governing permissions and
|
|
14
14
|
# limitations under the License.
|
|
15
15
|
|
|
16
|
-
|
|
16
|
+
# TODO: Remove after cleaning up the public API.
|
|
17
17
|
|
|
18
|
+
from warp._src.optim import adam as _adam
|
|
18
19
|
|
|
19
|
-
@wp.kernel
|
|
20
|
-
def adam_step_kernel_vec3(
|
|
21
|
-
g: wp.array(dtype=wp.vec3),
|
|
22
|
-
m: wp.array(dtype=wp.vec3),
|
|
23
|
-
v: wp.array(dtype=wp.vec3),
|
|
24
|
-
lr: float,
|
|
25
|
-
beta1: float,
|
|
26
|
-
beta2: float,
|
|
27
|
-
t: float,
|
|
28
|
-
eps: float,
|
|
29
|
-
params: wp.array(dtype=wp.vec3),
|
|
30
|
-
):
|
|
31
|
-
i = wp.tid()
|
|
32
|
-
m[i] = beta1 * m[i] + (1.0 - beta1) * g[i]
|
|
33
|
-
v[i] = beta2 * v[i] + (1.0 - beta2) * wp.cw_mul(g[i], g[i])
|
|
34
|
-
mhat = m[i] / (1.0 - wp.pow(beta1, (t + 1.0)))
|
|
35
|
-
vhat = v[i] / (1.0 - wp.pow(beta2, (t + 1.0)))
|
|
36
|
-
sqrt_vhat = wp.vec3(wp.sqrt(vhat[0]), wp.sqrt(vhat[1]), wp.sqrt(vhat[2]))
|
|
37
|
-
eps_vec3 = wp.vec3(eps, eps, eps)
|
|
38
|
-
params[i] = params[i] - lr * wp.cw_div(mhat, (sqrt_vhat + eps_vec3))
|
|
39
20
|
|
|
21
|
+
def __getattr__(name):
|
|
22
|
+
from warp._src.utils import get_deprecated_api
|
|
40
23
|
|
|
41
|
-
|
|
42
|
-
def adam_step_kernel_float(
|
|
43
|
-
g: wp.array(dtype=float),
|
|
44
|
-
m: wp.array(dtype=float),
|
|
45
|
-
v: wp.array(dtype=float),
|
|
46
|
-
lr: float,
|
|
47
|
-
beta1: float,
|
|
48
|
-
beta2: float,
|
|
49
|
-
t: float,
|
|
50
|
-
eps: float,
|
|
51
|
-
params: wp.array(dtype=float),
|
|
52
|
-
):
|
|
53
|
-
i = wp.tid()
|
|
54
|
-
m[i] = beta1 * m[i] + (1.0 - beta1) * g[i]
|
|
55
|
-
v[i] = beta2 * v[i] + (1.0 - beta2) * g[i] * g[i]
|
|
56
|
-
mhat = m[i] / (1.0 - wp.pow(beta1, (t + 1.0)))
|
|
57
|
-
vhat = v[i] / (1.0 - wp.pow(beta2, (t + 1.0)))
|
|
58
|
-
params[i] = params[i] - lr * mhat / (wp.sqrt(vhat) + eps)
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
@wp.kernel
|
|
62
|
-
def adam_step_kernel_half(
|
|
63
|
-
g: wp.array(dtype=wp.float16),
|
|
64
|
-
m: wp.array(dtype=float),
|
|
65
|
-
v: wp.array(dtype=float),
|
|
66
|
-
lr: float,
|
|
67
|
-
beta1: float,
|
|
68
|
-
beta2: float,
|
|
69
|
-
t: float,
|
|
70
|
-
eps: float,
|
|
71
|
-
params: wp.array(dtype=wp.float16),
|
|
72
|
-
):
|
|
73
|
-
i = wp.tid()
|
|
74
|
-
m[i] = beta1 * m[i] + (1.0 - beta1) * float(g[i])
|
|
75
|
-
v[i] = beta2 * v[i] + (1.0 - beta2) * float(g[i]) * float(g[i])
|
|
76
|
-
mhat = m[i] / (1.0 - wp.pow(beta1, (t + 1.0)))
|
|
77
|
-
vhat = v[i] / (1.0 - wp.pow(beta2, (t + 1.0)))
|
|
78
|
-
params[i] = params[i] - wp.float16(lr * mhat / (wp.sqrt(vhat) + eps))
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
class Adam:
|
|
82
|
-
"""An implementation of the Adam Optimizer
|
|
83
|
-
It is designed to mimic Pytorch's version.
|
|
84
|
-
https://pytorch.org/docs/stable/generated/torch.optim.Adam.html#torch.optim.Adam
|
|
85
|
-
"""
|
|
86
|
-
|
|
87
|
-
def __init__(self, params=None, lr=0.001, betas=(0.9, 0.999), eps=1e-08):
|
|
88
|
-
self.m = [] # first moment
|
|
89
|
-
self.v = [] # second moment
|
|
90
|
-
self.set_params(params)
|
|
91
|
-
self.lr = lr
|
|
92
|
-
self.beta1 = betas[0]
|
|
93
|
-
self.beta2 = betas[1]
|
|
94
|
-
self.eps = eps
|
|
95
|
-
self.t = 0
|
|
96
|
-
|
|
97
|
-
def set_params(self, params):
|
|
98
|
-
self.params = params
|
|
99
|
-
if params is not None and isinstance(params, list) and len(params) > 0:
|
|
100
|
-
if len(self.m) != len(params):
|
|
101
|
-
self.m = [None] * len(params) # reset first moment
|
|
102
|
-
if len(self.v) != len(params):
|
|
103
|
-
self.v = [None] * len(params) # reset second moment
|
|
104
|
-
for i in range(len(params)):
|
|
105
|
-
param = params[i]
|
|
106
|
-
|
|
107
|
-
if param.dtype == wp.vec3:
|
|
108
|
-
dtype = wp.vec3
|
|
109
|
-
elif param.dtype == wp.float32:
|
|
110
|
-
dtype = wp.float32
|
|
111
|
-
elif param.dtype == wp.float16:
|
|
112
|
-
dtype = wp.float32 # we always use fp32 for moments, even if params are fp16
|
|
113
|
-
else:
|
|
114
|
-
raise RuntimeError(f"Unsupported dtype for Warp Adam optimizer: {param.dtype}")
|
|
115
|
-
|
|
116
|
-
if self.m[i] is None or self.m[i].shape != param.shape or self.m[i].dtype != param.dtype:
|
|
117
|
-
self.m[i] = wp.zeros(shape=param.shape, dtype=dtype, device=param.device)
|
|
118
|
-
if self.v[i] is None or self.v[i].shape != param.shape or self.v[i].dtype != param.dtype:
|
|
119
|
-
self.v[i] = wp.zeros(shape=param.shape, dtype=dtype, device=param.device)
|
|
120
|
-
|
|
121
|
-
def reset_internal_state(self):
|
|
122
|
-
for m_i in self.m:
|
|
123
|
-
m_i.zero_()
|
|
124
|
-
for v_i in self.v:
|
|
125
|
-
v_i.zero_()
|
|
126
|
-
self.t = 0
|
|
127
|
-
|
|
128
|
-
def step(self, grad):
|
|
129
|
-
assert self.params is not None
|
|
130
|
-
for i in range(len(self.params)):
|
|
131
|
-
Adam.step_detail(
|
|
132
|
-
grad[i], self.m[i], self.v[i], self.lr, self.beta1, self.beta2, self.t, self.eps, self.params[i]
|
|
133
|
-
)
|
|
134
|
-
self.t = self.t + 1
|
|
135
|
-
|
|
136
|
-
@staticmethod
|
|
137
|
-
def step_detail(g, m, v, lr, beta1, beta2, t, eps, params):
|
|
138
|
-
assert params.dtype == g.dtype
|
|
139
|
-
assert params.shape == g.shape
|
|
140
|
-
kernel_inputs = [g, m, v, lr, beta1, beta2, t, eps, params]
|
|
141
|
-
if params.dtype == wp.types.float32:
|
|
142
|
-
wp.launch(
|
|
143
|
-
kernel=adam_step_kernel_float,
|
|
144
|
-
dim=len(params),
|
|
145
|
-
inputs=kernel_inputs,
|
|
146
|
-
device=params.device,
|
|
147
|
-
)
|
|
148
|
-
elif params.dtype == wp.types.float16:
|
|
149
|
-
wp.launch(
|
|
150
|
-
kernel=adam_step_kernel_half,
|
|
151
|
-
dim=len(params),
|
|
152
|
-
inputs=kernel_inputs,
|
|
153
|
-
device=params.device,
|
|
154
|
-
)
|
|
155
|
-
elif params.dtype == wp.types.vec3:
|
|
156
|
-
wp.launch(
|
|
157
|
-
kernel=adam_step_kernel_vec3,
|
|
158
|
-
dim=len(params),
|
|
159
|
-
inputs=kernel_inputs,
|
|
160
|
-
device=params.device,
|
|
161
|
-
)
|
|
162
|
-
else:
|
|
163
|
-
raise RuntimeError("Params data type not supported in Adam step kernels.")
|
|
24
|
+
return get_deprecated_api(_adam, "wp.optim", name)
|