warp-lang 1.0.1__py3-none-macosx_10_13_universal2.whl → 1.1.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.
- warp/__init__.py +108 -97
- warp/__init__.pyi +1 -1
- warp/bin/libwarp-clang.dylib +0 -0
- warp/bin/libwarp.dylib +0 -0
- warp/build.py +115 -113
- warp/build_dll.py +383 -375
- warp/builtins.py +3425 -3354
- warp/codegen.py +2878 -2792
- warp/config.py +40 -36
- warp/constants.py +45 -45
- warp/context.py +5194 -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 +383 -383
- warp/examples/benchmarks/benchmark_cloth.py +278 -279
- 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 +146 -146
- warp/examples/benchmarks/benchmark_launches.py +295 -295
- warp/examples/browse.py +29 -28
- warp/examples/core/example_dem.py +234 -221
- warp/examples/core/example_fluid.py +293 -267
- warp/examples/core/example_graph_capture.py +144 -129
- warp/examples/core/example_marching_cubes.py +188 -176
- warp/examples/core/example_mesh.py +174 -154
- warp/examples/core/example_mesh_intersect.py +205 -193
- warp/examples/core/example_nvdb.py +176 -169
- warp/examples/core/example_raycast.py +105 -89
- warp/examples/core/example_raymarch.py +199 -178
- warp/examples/core/example_render_opengl.py +185 -141
- warp/examples/core/example_sph.py +405 -389
- warp/examples/core/example_torch.py +222 -181
- warp/examples/core/example_wave.py +263 -249
- warp/examples/fem/bsr_utils.py +378 -380
- warp/examples/fem/example_apic_fluid.py +407 -391
- warp/examples/fem/example_convection_diffusion.py +182 -168
- warp/examples/fem/example_convection_diffusion_dg.py +219 -209
- warp/examples/fem/example_convection_diffusion_dg0.py +204 -194
- warp/examples/fem/example_deformed_geometry.py +177 -159
- warp/examples/fem/example_diffusion.py +201 -173
- warp/examples/fem/example_diffusion_3d.py +177 -152
- warp/examples/fem/example_diffusion_mgpu.py +221 -214
- warp/examples/fem/example_mixed_elasticity.py +244 -222
- warp/examples/fem/example_navier_stokes.py +259 -243
- warp/examples/fem/example_stokes.py +220 -192
- warp/examples/fem/example_stokes_transfer.py +265 -249
- warp/examples/fem/mesh_utils.py +133 -109
- warp/examples/fem/plot_utils.py +292 -287
- warp/examples/optim/example_bounce.py +260 -248
- warp/examples/optim/example_cloth_throw.py +222 -210
- warp/examples/optim/example_diffray.py +566 -535
- warp/examples/optim/example_drone.py +864 -835
- warp/examples/optim/example_inverse_kinematics.py +176 -169
- warp/examples/optim/example_inverse_kinematics_torch.py +185 -170
- warp/examples/optim/example_spring_cage.py +239 -234
- warp/examples/optim/example_trajectory.py +223 -201
- warp/examples/optim/example_walker.py +306 -292
- warp/examples/sim/example_cartpole.py +139 -128
- warp/examples/sim/example_cloth.py +196 -184
- warp/examples/sim/example_granular.py +124 -113
- warp/examples/sim/example_granular_collision_sdf.py +197 -185
- warp/examples/sim/example_jacobian_ik.py +236 -213
- warp/examples/sim/example_particle_chain.py +118 -106
- warp/examples/sim/example_quadruped.py +193 -179
- warp/examples/sim/example_rigid_chain.py +197 -189
- warp/examples/sim/example_rigid_contact.py +189 -176
- warp/examples/sim/example_rigid_force.py +127 -126
- warp/examples/sim/example_rigid_gyroscopic.py +109 -97
- warp/examples/sim/example_rigid_soft_contact.py +134 -124
- warp/examples/sim/example_soft_body.py +190 -178
- warp/fabric.py +337 -335
- warp/fem/__init__.py +60 -27
- warp/fem/cache.py +401 -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 +15 -19
- warp/fem/geometry/closest_point.py +69 -70
- warp/fem/geometry/deformed_geometry.py +270 -271
- warp/fem/geometry/element.py +744 -744
- warp/fem/geometry/geometry.py +184 -186
- warp/fem/geometry/grid_2d.py +380 -373
- warp/fem/geometry/grid_3d.py +441 -435
- warp/fem/geometry/hexmesh.py +953 -953
- 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 +1630 -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 +294 -292
- warp/fem/space/basis_space.py +488 -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 +267 -267
- warp/fem/space/grid_3d_function_space.py +305 -306
- warp/fem/space/hexmesh_function_space.py +350 -352
- warp/fem/space/partition.py +350 -350
- warp/fem/space/quadmesh_2d_function_space.py +368 -369
- warp/fem/space/restriction.py +158 -160
- warp/fem/space/shape/__init__.py +13 -15
- warp/fem/space/shape/cube_shape_function.py +738 -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 +294 -292
- warp/fem/space/topology.py +297 -295
- warp/fem/space/trimesh_2d_function_space.py +223 -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 +1072 -1025
- warp/native/builtin.h +1560 -1560
- warp/native/bvh.cpp +398 -398
- warp/native/bvh.cu +525 -525
- warp/native/bvh.h +429 -429
- warp/native/clang/clang.cpp +495 -464
- warp/native/crt.cpp +31 -31
- warp/native/crt.h +334 -334
- 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/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 +1498 -1498
- warp/native/matnn.h +333 -333
- warp/native/mesh.cpp +203 -203
- warp/native/mesh.cu +293 -293
- warp/native/mesh.h +1887 -1887
- warp/native/nanovdb/NanoVDB.h +4782 -4782
- warp/native/nanovdb/PNanoVDB.h +2553 -2553
- warp/native/nanovdb/PNanoVDBWrite.h +294 -294
- warp/native/noise.h +850 -850
- warp/native/quat.h +1084 -1084
- warp/native/rand.h +299 -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 +1132 -1132
- warp/native/volume.cpp +297 -297
- warp/native/volume.cu +32 -32
- warp/native/volume.h +538 -538
- warp/native/volume_builder.cu +425 -425
- warp/native/volume_builder.h +19 -19
- warp/native/warp.cpp +1057 -1052
- warp/native/warp.cu +2943 -2828
- warp/native/warp.h +313 -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 +3217 -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 +1910 -1991
- warp/sim/integrator_xpbd.py +3294 -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 +1227 -1227
- warp/stubs.py +2109 -2469
- warp/tape.py +1162 -225
- warp/tests/__init__.py +1 -1
- warp/tests/__main__.py +4 -4
- 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 +22 -22
- warp/tests/aux_test_grad_customs.py +23 -23
- warp/tests/aux_test_reference.py +11 -11
- warp/tests/aux_test_reference_reference.py +10 -10
- warp/tests/aux_test_square.py +17 -17
- warp/tests/aux_test_unresolved_func.py +14 -14
- warp/tests/aux_test_unresolved_symbol.py +14 -14
- warp/tests/disabled_kinematics.py +239 -239
- warp/tests/run_coverage_serial.py +31 -31
- warp/tests/test_adam.py +157 -157
- warp/tests/test_arithmetic.py +1124 -1124
- warp/tests/test_array.py +2417 -2326
- warp/tests/test_array_reduce.py +150 -150
- warp/tests/test_async.py +668 -656
- warp/tests/test_atomic.py +141 -141
- warp/tests/test_bool.py +204 -149
- warp/tests/test_builtins_resolution.py +1292 -1292
- warp/tests/test_bvh.py +164 -171
- warp/tests/test_closest_point_edge_edge.py +228 -228
- warp/tests/test_codegen.py +566 -553
- warp/tests/test_compile_consts.py +97 -101
- warp/tests/test_conditional.py +246 -246
- warp/tests/test_copy.py +232 -215
- warp/tests/test_ctypes.py +632 -632
- warp/tests/test_dense.py +67 -67
- warp/tests/test_devices.py +91 -98
- warp/tests/test_dlpack.py +530 -529
- warp/tests/test_examples.py +400 -378
- warp/tests/test_fabricarray.py +955 -955
- warp/tests/test_fast_math.py +62 -54
- warp/tests/test_fem.py +1277 -1278
- warp/tests/test_fp16.py +130 -130
- warp/tests/test_func.py +338 -337
- warp/tests/test_generics.py +571 -571
- warp/tests/test_grad.py +746 -640
- warp/tests/test_grad_customs.py +333 -336
- warp/tests/test_hash_grid.py +210 -164
- warp/tests/test_import.py +39 -39
- warp/tests/test_indexedarray.py +1134 -1134
- warp/tests/test_intersect.py +67 -67
- warp/tests/test_jax.py +307 -307
- warp/tests/test_large.py +167 -164
- warp/tests/test_launch.py +354 -354
- warp/tests/test_lerp.py +261 -261
- warp/tests/test_linear_solvers.py +191 -171
- warp/tests/test_lvalue.py +421 -493
- warp/tests/test_marching_cubes.py +65 -65
- warp/tests/test_mat.py +1801 -1827
- warp/tests/test_mat_lite.py +115 -115
- warp/tests/test_mat_scalar_ops.py +2907 -2889
- warp/tests/test_math.py +126 -193
- warp/tests/test_matmul.py +500 -499
- warp/tests/test_matmul_lite.py +410 -410
- warp/tests/test_mempool.py +188 -190
- warp/tests/test_mesh.py +284 -324
- warp/tests/test_mesh_query_aabb.py +228 -241
- warp/tests/test_mesh_query_point.py +692 -702
- warp/tests/test_mesh_query_ray.py +292 -303
- warp/tests/test_mlp.py +276 -276
- warp/tests/test_model.py +110 -110
- warp/tests/test_modules_lite.py +39 -39
- warp/tests/test_multigpu.py +163 -163
- warp/tests/test_noise.py +248 -248
- warp/tests/test_operators.py +250 -250
- warp/tests/test_options.py +123 -125
- warp/tests/test_peer.py +133 -137
- warp/tests/test_pinned.py +78 -78
- warp/tests/test_print.py +54 -54
- warp/tests/test_quat.py +2086 -2086
- warp/tests/test_rand.py +288 -288
- warp/tests/test_reload.py +217 -217
- warp/tests/test_rounding.py +179 -179
- warp/tests/test_runlength_encode.py +190 -190
- warp/tests/test_sim_grad.py +243 -0
- warp/tests/test_sim_kinematics.py +91 -97
- warp/tests/test_smoothstep.py +168 -168
- warp/tests/test_snippet.py +305 -266
- warp/tests/test_sparse.py +468 -460
- warp/tests/test_spatial.py +2148 -2148
- warp/tests/test_streams.py +486 -473
- warp/tests/test_struct.py +710 -675
- warp/tests/test_tape.py +173 -148
- warp/tests/test_torch.py +743 -743
- warp/tests/test_transient_module.py +87 -87
- warp/tests/test_types.py +556 -659
- warp/tests/test_utils.py +490 -499
- warp/tests/test_vec.py +1264 -1268
- warp/tests/test_vec_lite.py +73 -73
- warp/tests/test_vec_scalar_ops.py +2099 -2099
- warp/tests/test_verify_fp.py +94 -94
- warp/tests/test_volume.py +737 -736
- warp/tests/test_volume_write.py +255 -265
- warp/tests/unittest_serial.py +37 -37
- warp/tests/unittest_suites.py +363 -359
- warp/tests/unittest_utils.py +603 -578
- warp/tests/unused_test_misc.py +71 -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 +566 -561
- warp/torch.py +321 -295
- warp/types.py +4504 -4450
- warp/utils.py +1008 -821
- {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/LICENSE.md +126 -126
- {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/METADATA +338 -400
- warp_lang-1.1.0.dist-info/RECORD +352 -0
- warp/examples/assets/cube.usda +0 -42
- warp/examples/assets/sphere.usda +0 -56
- warp/examples/assets/torus.usda +0 -105
- warp_lang-1.0.1.dist-info/RECORD +0 -352
- {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/WHEEL +0 -0
- {warp_lang-1.0.1.dist-info → warp_lang-1.1.0.dist-info}/top_level.txt +0 -0
warp/native/cuda_util.h
CHANGED
|
@@ -1,203 +1,288 @@
|
|
|
1
|
-
/** Copyright (c) 2022 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
|
-
|
|
9
|
-
#pragma once
|
|
10
|
-
|
|
11
|
-
#include "builtin.h"
|
|
12
|
-
|
|
13
|
-
#if WP_ENABLE_CUDA
|
|
14
|
-
|
|
15
|
-
#include <cudaTypedefs.h>
|
|
16
|
-
#include <cuda_runtime_api.h>
|
|
17
|
-
|
|
18
|
-
#include <stdio.h>
|
|
19
|
-
|
|
20
|
-
#include <vector>
|
|
21
|
-
|
|
22
|
-
#define check_cuda(code) (check_cuda_result(code, __FUNCTION__, __FILE__, __LINE__))
|
|
23
|
-
#define check_cu(code) (check_cu_result(code, __FUNCTION__, __FILE__, __LINE__))
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
#if defined(__CUDACC__)
|
|
27
|
-
#if _DEBUG
|
|
28
|
-
// helper for launching kernels (synchronize + error checking after each kernel)
|
|
29
|
-
#define wp_launch_device(context, kernel, dim, args) { \
|
|
30
|
-
if (dim) { \
|
|
31
|
-
ContextGuard guard(context); \
|
|
32
|
-
|
|
33
|
-
const int
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
CUresult
|
|
55
|
-
CUresult
|
|
56
|
-
CUresult
|
|
57
|
-
CUresult
|
|
58
|
-
CUresult
|
|
59
|
-
CUresult
|
|
60
|
-
CUresult
|
|
61
|
-
CUresult
|
|
62
|
-
CUresult
|
|
63
|
-
CUresult
|
|
64
|
-
CUresult
|
|
65
|
-
CUresult
|
|
66
|
-
CUresult
|
|
67
|
-
CUresult
|
|
68
|
-
CUresult
|
|
69
|
-
CUresult
|
|
70
|
-
CUresult
|
|
71
|
-
CUresult
|
|
72
|
-
CUresult
|
|
73
|
-
CUresult
|
|
74
|
-
CUresult
|
|
75
|
-
CUresult
|
|
76
|
-
CUresult
|
|
77
|
-
CUresult
|
|
78
|
-
CUresult
|
|
79
|
-
CUresult
|
|
80
|
-
CUresult
|
|
81
|
-
CUresult
|
|
82
|
-
CUresult
|
|
83
|
-
CUresult
|
|
84
|
-
CUresult
|
|
85
|
-
CUresult
|
|
86
|
-
CUresult
|
|
87
|
-
CUresult
|
|
88
|
-
CUresult
|
|
89
|
-
CUresult
|
|
90
|
-
CUresult
|
|
91
|
-
CUresult
|
|
92
|
-
CUresult
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
bool
|
|
106
|
-
|
|
107
|
-
inline
|
|
108
|
-
{
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
171
|
-
|
|
172
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
{
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
(
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
|
|
1
|
+
/** Copyright (c) 2022 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
|
+
|
|
9
|
+
#pragma once
|
|
10
|
+
|
|
11
|
+
#include "builtin.h"
|
|
12
|
+
|
|
13
|
+
#if WP_ENABLE_CUDA
|
|
14
|
+
|
|
15
|
+
#include <cudaTypedefs.h>
|
|
16
|
+
#include <cuda_runtime_api.h>
|
|
17
|
+
|
|
18
|
+
#include <stdio.h>
|
|
19
|
+
|
|
20
|
+
#include <vector>
|
|
21
|
+
|
|
22
|
+
#define check_cuda(code) (check_cuda_result(code, __FUNCTION__, __FILE__, __LINE__))
|
|
23
|
+
#define check_cu(code) (check_cu_result(code, __FUNCTION__, __FILE__, __LINE__))
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
#if defined(__CUDACC__)
|
|
27
|
+
#if _DEBUG
|
|
28
|
+
// helper for launching kernels (synchronize + error checking after each kernel)
|
|
29
|
+
#define wp_launch_device(context, kernel, dim, args) { \
|
|
30
|
+
if (dim) { \
|
|
31
|
+
ContextGuard guard(context); \
|
|
32
|
+
cudaStream_t stream = (cudaStream_t)cuda_stream_get_current(); \
|
|
33
|
+
const int num_threads = 256; \
|
|
34
|
+
const int num_blocks = (dim+num_threads-1)/num_threads; \
|
|
35
|
+
begin_cuda_range(WP_TIMING_KERNEL_BUILTIN, stream, context, #kernel); \
|
|
36
|
+
kernel<<<num_blocks, 256, 0, stream>>>args; \
|
|
37
|
+
check_cuda(cuda_context_check(WP_CURRENT_CONTEXT)); \
|
|
38
|
+
end_cuda_range(WP_TIMING_KERNEL_BUILTIN, stream); }}
|
|
39
|
+
#else
|
|
40
|
+
// helper for launching kernels (no error checking)
|
|
41
|
+
#define wp_launch_device(context, kernel, dim, args) { \
|
|
42
|
+
if (dim) { \
|
|
43
|
+
ContextGuard guard(context); \
|
|
44
|
+
cudaStream_t stream = (cudaStream_t)cuda_stream_get_current(); \
|
|
45
|
+
const int num_threads = 256; \
|
|
46
|
+
const int num_blocks = (dim+num_threads-1)/num_threads; \
|
|
47
|
+
begin_cuda_range(WP_TIMING_KERNEL_BUILTIN, stream, context, #kernel); \
|
|
48
|
+
kernel<<<num_blocks, 256, 0, stream>>>args; \
|
|
49
|
+
end_cuda_range(WP_TIMING_KERNEL_BUILTIN, stream); }}
|
|
50
|
+
#endif // _DEBUG
|
|
51
|
+
#endif // defined(__CUDACC__)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
CUresult cuDriverGetVersion_f(int* version);
|
|
55
|
+
CUresult cuGetErrorName_f(CUresult result, const char** pstr);
|
|
56
|
+
CUresult cuGetErrorString_f(CUresult result, const char** pstr);
|
|
57
|
+
CUresult cuInit_f(unsigned int flags);
|
|
58
|
+
CUresult cuDeviceGet_f(CUdevice *dev, int ordinal);
|
|
59
|
+
CUresult cuDeviceGetCount_f(int* count);
|
|
60
|
+
CUresult cuDeviceGetName_f(char* name, int len, CUdevice dev);
|
|
61
|
+
CUresult cuDeviceGetAttribute_f(int* value, CUdevice_attribute attrib, CUdevice dev);
|
|
62
|
+
CUresult cuDeviceGetUuid_f(CUuuid* uuid, CUdevice dev);
|
|
63
|
+
CUresult cuDevicePrimaryCtxRetain_f(CUcontext* ctx, CUdevice dev);
|
|
64
|
+
CUresult cuDevicePrimaryCtxRelease_f(CUdevice dev);
|
|
65
|
+
CUresult cuDeviceCanAccessPeer_f(int* can_access, CUdevice dev, CUdevice peer_dev);
|
|
66
|
+
CUresult cuMemGetInfo_f(size_t* free, size_t* total);
|
|
67
|
+
CUresult cuCtxGetCurrent_f(CUcontext* ctx);
|
|
68
|
+
CUresult cuCtxSetCurrent_f(CUcontext ctx);
|
|
69
|
+
CUresult cuCtxPushCurrent_f(CUcontext ctx);
|
|
70
|
+
CUresult cuCtxPopCurrent_f(CUcontext* ctx);
|
|
71
|
+
CUresult cuCtxSynchronize_f();
|
|
72
|
+
CUresult cuCtxGetDevice_f(CUdevice* dev);
|
|
73
|
+
CUresult cuCtxCreate_f(CUcontext* ctx, unsigned int flags, CUdevice dev);
|
|
74
|
+
CUresult cuCtxDestroy_f(CUcontext ctx);
|
|
75
|
+
CUresult cuCtxEnablePeerAccess_f(CUcontext peer_ctx, unsigned int flags);
|
|
76
|
+
CUresult cuCtxDisablePeerAccess_f(CUcontext peer_ctx);
|
|
77
|
+
CUresult cuStreamCreate_f(CUstream* stream, unsigned int flags);
|
|
78
|
+
CUresult cuStreamDestroy_f(CUstream stream);
|
|
79
|
+
CUresult cuStreamSynchronize_f(CUstream stream);
|
|
80
|
+
CUresult cuStreamWaitEvent_f(CUstream stream, CUevent event, unsigned int flags);
|
|
81
|
+
CUresult cuStreamGetCtx_f(CUstream stream, CUcontext* pctx);
|
|
82
|
+
CUresult cuStreamGetCaptureInfo_f(CUstream stream, CUstreamCaptureStatus *captureStatus_out, cuuint64_t *id_out, CUgraph *graph_out, const CUgraphNode **dependencies_out, size_t *numDependencies_out);
|
|
83
|
+
CUresult cuStreamUpdateCaptureDependencies_f(CUstream stream, CUgraphNode *dependencies, size_t numDependencies, unsigned int flags);
|
|
84
|
+
CUresult cuEventCreate_f(CUevent* event, unsigned int flags);
|
|
85
|
+
CUresult cuEventDestroy_f(CUevent event);
|
|
86
|
+
CUresult cuEventRecord_f(CUevent event, CUstream stream);
|
|
87
|
+
CUresult cuEventRecordWithFlags_f(CUevent event, CUstream stream, unsigned int flags);
|
|
88
|
+
CUresult cuEventSynchronize_f(CUevent event);
|
|
89
|
+
CUresult cuModuleUnload_f(CUmodule hmod);
|
|
90
|
+
CUresult cuModuleLoadDataEx_f(CUmodule *module, const void *image, unsigned int numOptions, CUjit_option *options, void **optionValues);
|
|
91
|
+
CUresult cuModuleGetFunction_f(CUfunction *hfunc, CUmodule hmod, const char *name);
|
|
92
|
+
CUresult cuLaunchKernel_f(CUfunction f, unsigned int gridDimX, unsigned int gridDimY, unsigned int gridDimZ, unsigned int blockDimX, unsigned int blockDimY, unsigned int blockDimZ, unsigned int sharedMemBytes, CUstream hStream, void **kernelParams, void **extra);
|
|
93
|
+
CUresult cuMemcpyPeerAsync_f(CUdeviceptr dst_ptr, CUcontext dst_ctx, CUdeviceptr src_ptr, CUcontext src_ctx, size_t n, CUstream stream);
|
|
94
|
+
CUresult cuPointerGetAttribute_f(void* data, CUpointer_attribute attribute, CUdeviceptr ptr);
|
|
95
|
+
CUresult cuGraphicsMapResources_f(unsigned int count, CUgraphicsResource* resources, CUstream stream);
|
|
96
|
+
CUresult cuGraphicsUnmapResources_f(unsigned int count, CUgraphicsResource* resources, CUstream hStream);
|
|
97
|
+
CUresult cuGraphicsResourceGetMappedPointer_f(CUdeviceptr* pDevPtr, size_t* pSize, CUgraphicsResource resource);
|
|
98
|
+
CUresult cuGraphicsGLRegisterBuffer_f(CUgraphicsResource *pCudaResource, unsigned int buffer, unsigned int flags);
|
|
99
|
+
CUresult cuGraphicsUnregisterResource_f(CUgraphicsResource resource);
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
bool init_cuda_driver();
|
|
103
|
+
bool is_cuda_driver_initialized();
|
|
104
|
+
|
|
105
|
+
bool check_cuda_result(cudaError_t code, const char* func, const char* file, int line);
|
|
106
|
+
|
|
107
|
+
inline bool check_cuda_result(uint64_t code, const char* func, const char* file, int line)
|
|
108
|
+
{
|
|
109
|
+
return check_cuda_result(static_cast<cudaError_t>(code), func, file, line);
|
|
110
|
+
}
|
|
111
|
+
|
|
112
|
+
bool check_cu_result(CUresult result, const char* func, const char* file, int line);
|
|
113
|
+
|
|
114
|
+
inline uint64_t get_capture_id(CUstream stream)
|
|
115
|
+
{
|
|
116
|
+
CUstreamCaptureStatus status;
|
|
117
|
+
uint64_t id = 0;
|
|
118
|
+
check_cu(cuStreamGetCaptureInfo_f(stream, &status, &id, NULL, NULL, NULL));
|
|
119
|
+
return id;
|
|
120
|
+
}
|
|
121
|
+
|
|
122
|
+
inline CUgraph get_capture_graph(CUstream stream)
|
|
123
|
+
{
|
|
124
|
+
CUstreamCaptureStatus status;
|
|
125
|
+
CUgraph graph = NULL;
|
|
126
|
+
check_cu(cuStreamGetCaptureInfo_f(stream, &status, NULL, &graph, NULL, NULL));
|
|
127
|
+
return graph;
|
|
128
|
+
}
|
|
129
|
+
|
|
130
|
+
bool get_capture_dependencies(CUstream stream, std::vector<CUgraphNode>& dependencies_ret);
|
|
131
|
+
|
|
132
|
+
bool get_graph_leaf_nodes(cudaGraph_t graph, std::vector<cudaGraphNode_t>& leaf_nodes_ret);
|
|
133
|
+
|
|
134
|
+
inline CUcontext get_stream_context(CUstream stream)
|
|
135
|
+
{
|
|
136
|
+
CUcontext context;
|
|
137
|
+
if (check_cu(cuStreamGetCtx_f(stream, &context)))
|
|
138
|
+
return context;
|
|
139
|
+
else
|
|
140
|
+
return NULL;
|
|
141
|
+
}
|
|
142
|
+
|
|
143
|
+
inline CUcontext get_stream_context(void* stream)
|
|
144
|
+
{
|
|
145
|
+
return get_stream_context(static_cast<CUstream>(stream));
|
|
146
|
+
}
|
|
147
|
+
|
|
148
|
+
|
|
149
|
+
//
|
|
150
|
+
// Scoped CUDA context guard
|
|
151
|
+
//
|
|
152
|
+
// Behaviour on entry
|
|
153
|
+
// - If the given `context` is NULL, do nothing.
|
|
154
|
+
// - If the given `context` is the same as the current context, do nothing.
|
|
155
|
+
// - If the given `context` is different from the current context, make the given context current.
|
|
156
|
+
//
|
|
157
|
+
// Behaviour on exit
|
|
158
|
+
// - If the current context did not change on entry, do nothing.
|
|
159
|
+
// - If the `restore` flag was true on entry, make the previous context current.
|
|
160
|
+
//
|
|
161
|
+
// Default exit behaviour policy
|
|
162
|
+
// - If the `restore` flag is omitted on entry, fall back on the global `always_restore` flag.
|
|
163
|
+
// - This allows us to easily change the default behaviour of the guards.
|
|
164
|
+
//
|
|
165
|
+
class ContextGuard
|
|
166
|
+
{
|
|
167
|
+
public:
|
|
168
|
+
// default policy for restoring contexts
|
|
169
|
+
static bool always_restore;
|
|
170
|
+
|
|
171
|
+
explicit ContextGuard(CUcontext context, bool restore=always_restore)
|
|
172
|
+
: need_restore(false)
|
|
173
|
+
{
|
|
174
|
+
if (context)
|
|
175
|
+
{
|
|
176
|
+
if (check_cu(cuCtxGetCurrent_f(&prev_context)) && context != prev_context)
|
|
177
|
+
need_restore = check_cu(cuCtxSetCurrent_f(context)) && restore;
|
|
178
|
+
}
|
|
179
|
+
}
|
|
180
|
+
|
|
181
|
+
explicit ContextGuard(void* context, bool restore=always_restore)
|
|
182
|
+
: ContextGuard(static_cast<CUcontext>(context), restore)
|
|
183
|
+
{
|
|
184
|
+
}
|
|
185
|
+
|
|
186
|
+
~ContextGuard()
|
|
187
|
+
{
|
|
188
|
+
if (need_restore)
|
|
189
|
+
check_cu(cuCtxSetCurrent_f(prev_context));
|
|
190
|
+
}
|
|
191
|
+
|
|
192
|
+
private:
|
|
193
|
+
CUcontext prev_context;
|
|
194
|
+
bool need_restore;
|
|
195
|
+
};
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
// CUDA timing range used during event-based timing
|
|
199
|
+
struct CudaTimingRange
|
|
200
|
+
{
|
|
201
|
+
void* context;
|
|
202
|
+
const char* name;
|
|
203
|
+
int flag;
|
|
204
|
+
CUevent start;
|
|
205
|
+
CUevent end;
|
|
206
|
+
};
|
|
207
|
+
|
|
208
|
+
// Timing result used to pass timings to Python
|
|
209
|
+
struct timing_result_t
|
|
210
|
+
{
|
|
211
|
+
void* context;
|
|
212
|
+
const char* name;
|
|
213
|
+
int flag;
|
|
214
|
+
float elapsed;
|
|
215
|
+
};
|
|
216
|
+
|
|
217
|
+
struct CudaTimingState
|
|
218
|
+
{
|
|
219
|
+
int flags;
|
|
220
|
+
std::vector<CudaTimingRange> ranges;
|
|
221
|
+
CudaTimingState* parent;
|
|
222
|
+
|
|
223
|
+
CudaTimingState(int flags, CudaTimingState* parent)
|
|
224
|
+
: flags(flags), parent(parent)
|
|
225
|
+
{
|
|
226
|
+
}
|
|
227
|
+
};
|
|
228
|
+
|
|
229
|
+
// timing flags
|
|
230
|
+
constexpr int WP_TIMING_KERNEL = 1; // Warp kernel
|
|
231
|
+
constexpr int WP_TIMING_KERNEL_BUILTIN = 2; // internal kernel
|
|
232
|
+
constexpr int WP_TIMING_MEMCPY = 4; // memcpy operation
|
|
233
|
+
constexpr int WP_TIMING_MEMSET = 8; // memset operation
|
|
234
|
+
constexpr int WP_TIMING_GRAPH = 16; // graph launch
|
|
235
|
+
|
|
236
|
+
#define begin_cuda_range(_flag, _stream, _context, _name) \
|
|
237
|
+
CudaTimingRange _timing_range; \
|
|
238
|
+
bool _timing_enabled; \
|
|
239
|
+
if ((g_cuda_timing_state->flags & _flag) && !cuda_stream_is_capturing(_stream)) { \
|
|
240
|
+
ContextGuard guard(_context, true); \
|
|
241
|
+
_timing_enabled = true; \
|
|
242
|
+
_timing_range.context = _context ? _context : get_current_context(); \
|
|
243
|
+
_timing_range.name = _name; \
|
|
244
|
+
_timing_range.flag = _flag; \
|
|
245
|
+
check_cu(cuEventCreate_f(&_timing_range.start, CU_EVENT_DEFAULT)); \
|
|
246
|
+
check_cu(cuEventCreate_f(&_timing_range.end, CU_EVENT_DEFAULT)); \
|
|
247
|
+
check_cu(cuEventRecord_f(_timing_range.start, static_cast<CUstream>(_stream))); \
|
|
248
|
+
} else { \
|
|
249
|
+
_timing_enabled = false; \
|
|
250
|
+
}
|
|
251
|
+
|
|
252
|
+
#define end_cuda_range(_flag, _stream) \
|
|
253
|
+
if (_timing_enabled) { \
|
|
254
|
+
check_cu(cuEventRecord_f(_timing_range.end, static_cast<CUstream>(_stream))); \
|
|
255
|
+
g_cuda_timing_state->ranges.push_back(_timing_range); \
|
|
256
|
+
}
|
|
257
|
+
|
|
258
|
+
extern CudaTimingState* g_cuda_timing_state;
|
|
259
|
+
|
|
260
|
+
|
|
261
|
+
#else
|
|
262
|
+
|
|
263
|
+
typedef int CUdevice;
|
|
264
|
+
typedef struct CUctx_st* CUcontext;
|
|
265
|
+
typedef struct CUstream_st* CUstream;
|
|
266
|
+
|
|
267
|
+
class ContextGuard
|
|
268
|
+
{
|
|
269
|
+
public:
|
|
270
|
+
explicit ContextGuard(CUcontext context, bool restore=false)
|
|
271
|
+
{
|
|
272
|
+
(void)context;
|
|
273
|
+
(void)restore;
|
|
274
|
+
}
|
|
275
|
+
|
|
276
|
+
explicit ContextGuard(void* context, bool restore=false)
|
|
277
|
+
{
|
|
278
|
+
(void)context;
|
|
279
|
+
(void)restore;
|
|
280
|
+
}
|
|
281
|
+
};
|
|
282
|
+
|
|
283
|
+
#endif // WP_ENABLE_CUDA
|
|
284
|
+
|
|
285
|
+
// Pass this value to device functions as the `context` parameter to bypass unnecessary context management.
|
|
286
|
+
// This works in conjunction with ContextGuards, which do nothing if the given context is NULL.
|
|
287
|
+
// Using this variable instead of passing NULL directly aids readability and makes the intent clear.
|
|
288
|
+
constexpr void* WP_CURRENT_CONTEXT = NULL;
|
warp/native/cutlass_gemm.cpp
CHANGED
|
@@ -1,34 +1,34 @@
|
|
|
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
|
-
|
|
9
|
-
#include "builtin.h"
|
|
10
|
-
|
|
11
|
-
// stubs for platforms where there is no CUDA
|
|
12
|
-
#if !WP_ENABLE_CUDA || !WP_ENABLE_CUTLASS
|
|
13
|
-
|
|
14
|
-
extern "C"
|
|
15
|
-
{
|
|
16
|
-
|
|
17
|
-
WP_API
|
|
18
|
-
bool cutlass_gemm(
|
|
19
|
-
void* context, int compute_capability,
|
|
20
|
-
int m, int n, int k,
|
|
21
|
-
const char* datatype_str,
|
|
22
|
-
const void* a, const void* b, const void* c, void* d,
|
|
23
|
-
float alpha, float beta,
|
|
24
|
-
bool row_major_a, bool row_major_b,
|
|
25
|
-
bool allow_tf32x3_arith,
|
|
26
|
-
int batch_count)
|
|
27
|
-
{
|
|
28
|
-
printf("CUDA is disabled and/or CUTLASS is disabled.\n");
|
|
29
|
-
return false;
|
|
30
|
-
}
|
|
31
|
-
|
|
32
|
-
} // extern "C"
|
|
33
|
-
|
|
34
|
-
#endif // !WP_ENABLE_CUDA || !WP_ENABLE_CUTLASS
|
|
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
|
+
|
|
9
|
+
#include "builtin.h"
|
|
10
|
+
|
|
11
|
+
// stubs for platforms where there is no CUDA
|
|
12
|
+
#if !WP_ENABLE_CUDA || !WP_ENABLE_CUTLASS
|
|
13
|
+
|
|
14
|
+
extern "C"
|
|
15
|
+
{
|
|
16
|
+
|
|
17
|
+
WP_API
|
|
18
|
+
bool cutlass_gemm(
|
|
19
|
+
void* context, int compute_capability,
|
|
20
|
+
int m, int n, int k,
|
|
21
|
+
const char* datatype_str,
|
|
22
|
+
const void* a, const void* b, const void* c, void* d,
|
|
23
|
+
float alpha, float beta,
|
|
24
|
+
bool row_major_a, bool row_major_b,
|
|
25
|
+
bool allow_tf32x3_arith,
|
|
26
|
+
int batch_count)
|
|
27
|
+
{
|
|
28
|
+
printf("CUDA is disabled and/or CUTLASS is disabled.\n");
|
|
29
|
+
return false;
|
|
30
|
+
}
|
|
31
|
+
|
|
32
|
+
} // extern "C"
|
|
33
|
+
|
|
34
|
+
#endif // !WP_ENABLE_CUDA || !WP_ENABLE_CUTLASS
|