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.

Files changed (350) hide show
  1. warp/__init__.py +301 -287
  2. warp/__init__.pyi +2220 -313
  3. warp/_src/__init__.py +14 -0
  4. warp/_src/autograd.py +1075 -0
  5. warp/_src/build.py +618 -0
  6. warp/_src/build_dll.py +640 -0
  7. warp/{builtins.py → _src/builtins.py} +1497 -226
  8. warp/_src/codegen.py +4359 -0
  9. warp/{config.py → _src/config.py} +178 -169
  10. warp/_src/constants.py +57 -0
  11. warp/_src/context.py +8294 -0
  12. warp/_src/dlpack.py +462 -0
  13. warp/_src/fabric.py +355 -0
  14. warp/_src/fem/__init__.py +14 -0
  15. warp/_src/fem/adaptivity.py +508 -0
  16. warp/_src/fem/cache.py +687 -0
  17. warp/_src/fem/dirichlet.py +188 -0
  18. warp/{fem → _src/fem}/domain.py +40 -30
  19. warp/_src/fem/field/__init__.py +131 -0
  20. warp/_src/fem/field/field.py +701 -0
  21. warp/{fem → _src/fem}/field/nodal_field.py +30 -15
  22. warp/{fem → _src/fem}/field/restriction.py +1 -1
  23. warp/{fem → _src/fem}/field/virtual.py +53 -27
  24. warp/_src/fem/geometry/__init__.py +32 -0
  25. warp/{fem → _src/fem}/geometry/adaptive_nanogrid.py +77 -163
  26. warp/_src/fem/geometry/closest_point.py +97 -0
  27. warp/{fem → _src/fem}/geometry/deformed_geometry.py +14 -22
  28. warp/{fem → _src/fem}/geometry/element.py +32 -10
  29. warp/{fem → _src/fem}/geometry/geometry.py +48 -20
  30. warp/{fem → _src/fem}/geometry/grid_2d.py +12 -23
  31. warp/{fem → _src/fem}/geometry/grid_3d.py +12 -23
  32. warp/{fem → _src/fem}/geometry/hexmesh.py +40 -63
  33. warp/{fem → _src/fem}/geometry/nanogrid.py +255 -248
  34. warp/{fem → _src/fem}/geometry/partition.py +121 -63
  35. warp/{fem → _src/fem}/geometry/quadmesh.py +26 -45
  36. warp/{fem → _src/fem}/geometry/tetmesh.py +40 -63
  37. warp/{fem → _src/fem}/geometry/trimesh.py +26 -45
  38. warp/{fem → _src/fem}/integrate.py +164 -158
  39. warp/_src/fem/linalg.py +383 -0
  40. warp/_src/fem/operator.py +396 -0
  41. warp/_src/fem/polynomial.py +229 -0
  42. warp/{fem → _src/fem}/quadrature/pic_quadrature.py +15 -20
  43. warp/{fem → _src/fem}/quadrature/quadrature.py +95 -47
  44. warp/_src/fem/space/__init__.py +248 -0
  45. warp/{fem → _src/fem}/space/basis_function_space.py +20 -11
  46. warp/_src/fem/space/basis_space.py +679 -0
  47. warp/{fem → _src/fem}/space/dof_mapper.py +3 -3
  48. warp/{fem → _src/fem}/space/function_space.py +14 -13
  49. warp/{fem → _src/fem}/space/grid_2d_function_space.py +4 -7
  50. warp/{fem → _src/fem}/space/grid_3d_function_space.py +4 -4
  51. warp/{fem → _src/fem}/space/hexmesh_function_space.py +4 -10
  52. warp/{fem → _src/fem}/space/nanogrid_function_space.py +3 -9
  53. warp/{fem → _src/fem}/space/partition.py +117 -60
  54. warp/{fem → _src/fem}/space/quadmesh_function_space.py +4 -10
  55. warp/{fem → _src/fem}/space/restriction.py +66 -33
  56. warp/_src/fem/space/shape/__init__.py +152 -0
  57. warp/{fem → _src/fem}/space/shape/cube_shape_function.py +9 -9
  58. warp/{fem → _src/fem}/space/shape/shape_function.py +8 -9
  59. warp/{fem → _src/fem}/space/shape/square_shape_function.py +6 -6
  60. warp/{fem → _src/fem}/space/shape/tet_shape_function.py +3 -3
  61. warp/{fem → _src/fem}/space/shape/triangle_shape_function.py +3 -3
  62. warp/{fem → _src/fem}/space/tetmesh_function_space.py +3 -9
  63. warp/_src/fem/space/topology.py +459 -0
  64. warp/{fem → _src/fem}/space/trimesh_function_space.py +3 -9
  65. warp/_src/fem/types.py +112 -0
  66. warp/_src/fem/utils.py +486 -0
  67. warp/_src/jax.py +186 -0
  68. warp/_src/jax_experimental/__init__.py +14 -0
  69. warp/_src/jax_experimental/custom_call.py +387 -0
  70. warp/_src/jax_experimental/ffi.py +1284 -0
  71. warp/_src/jax_experimental/xla_ffi.py +656 -0
  72. warp/_src/marching_cubes.py +708 -0
  73. warp/_src/math.py +414 -0
  74. warp/_src/optim/__init__.py +14 -0
  75. warp/_src/optim/adam.py +163 -0
  76. warp/_src/optim/linear.py +1606 -0
  77. warp/_src/optim/sgd.py +112 -0
  78. warp/_src/paddle.py +406 -0
  79. warp/_src/render/__init__.py +14 -0
  80. warp/_src/render/imgui_manager.py +289 -0
  81. warp/_src/render/render_opengl.py +3636 -0
  82. warp/_src/render/render_usd.py +937 -0
  83. warp/_src/render/utils.py +160 -0
  84. warp/_src/sparse.py +2716 -0
  85. warp/_src/tape.py +1206 -0
  86. warp/{thirdparty → _src/thirdparty}/unittest_parallel.py +9 -2
  87. warp/_src/torch.py +391 -0
  88. warp/_src/types.py +5870 -0
  89. warp/_src/utils.py +1693 -0
  90. warp/autograd.py +12 -1054
  91. warp/bin/warp-clang.dll +0 -0
  92. warp/bin/warp.dll +0 -0
  93. warp/build.py +8 -588
  94. warp/build_dll.py +6 -471
  95. warp/codegen.py +6 -4246
  96. warp/constants.py +6 -39
  97. warp/context.py +12 -7851
  98. warp/dlpack.py +6 -444
  99. warp/examples/distributed/example_jacobi_mpi.py +4 -5
  100. warp/examples/fem/example_adaptive_grid.py +1 -1
  101. warp/examples/fem/example_apic_fluid.py +1 -1
  102. warp/examples/fem/example_burgers.py +8 -8
  103. warp/examples/fem/example_diffusion.py +1 -1
  104. warp/examples/fem/example_distortion_energy.py +1 -1
  105. warp/examples/fem/example_mixed_elasticity.py +2 -2
  106. warp/examples/fem/example_navier_stokes.py +1 -1
  107. warp/examples/fem/example_nonconforming_contact.py +7 -7
  108. warp/examples/fem/example_stokes.py +1 -1
  109. warp/examples/fem/example_stokes_transfer.py +1 -1
  110. warp/examples/fem/utils.py +2 -2
  111. warp/examples/interop/example_jax_callable.py +1 -1
  112. warp/examples/interop/example_jax_ffi_callback.py +1 -1
  113. warp/examples/interop/example_jax_kernel.py +3 -2
  114. warp/examples/tile/example_tile_mcgp.py +191 -0
  115. warp/fabric.py +6 -337
  116. warp/fem/__init__.py +159 -97
  117. warp/fem/adaptivity.py +7 -489
  118. warp/fem/cache.py +9 -648
  119. warp/fem/dirichlet.py +6 -184
  120. warp/fem/field/__init__.py +8 -109
  121. warp/fem/field/field.py +7 -652
  122. warp/fem/geometry/__init__.py +7 -18
  123. warp/fem/geometry/closest_point.py +11 -77
  124. warp/fem/linalg.py +18 -366
  125. warp/fem/operator.py +11 -369
  126. warp/fem/polynomial.py +9 -209
  127. warp/fem/space/__init__.py +5 -211
  128. warp/fem/space/basis_space.py +6 -662
  129. warp/fem/space/shape/__init__.py +41 -118
  130. warp/fem/space/topology.py +6 -437
  131. warp/fem/types.py +6 -81
  132. warp/fem/utils.py +11 -444
  133. warp/jax.py +8 -165
  134. warp/jax_experimental/__init__.py +14 -1
  135. warp/jax_experimental/custom_call.py +8 -342
  136. warp/jax_experimental/ffi.py +17 -853
  137. warp/jax_experimental/xla_ffi.py +5 -596
  138. warp/marching_cubes.py +5 -689
  139. warp/math.py +16 -393
  140. warp/native/array.h +385 -37
  141. warp/native/builtin.h +316 -39
  142. warp/native/bvh.cpp +43 -9
  143. warp/native/bvh.cu +62 -27
  144. warp/native/bvh.h +310 -309
  145. warp/native/clang/clang.cpp +102 -97
  146. warp/native/coloring.cpp +0 -1
  147. warp/native/crt.h +208 -0
  148. warp/native/exports.h +156 -0
  149. warp/native/hashgrid.cu +2 -0
  150. warp/native/intersect.h +24 -1
  151. warp/native/intersect_tri.h +44 -35
  152. warp/native/mat.h +1456 -276
  153. warp/native/mesh.cpp +4 -4
  154. warp/native/mesh.cu +4 -2
  155. warp/native/mesh.h +176 -61
  156. warp/native/quat.h +0 -52
  157. warp/native/scan.cu +2 -0
  158. warp/native/sort.cu +22 -13
  159. warp/native/sort.h +2 -0
  160. warp/native/sparse.cu +7 -3
  161. warp/native/spatial.h +12 -0
  162. warp/native/tile.h +837 -70
  163. warp/native/tile_radix_sort.h +1 -1
  164. warp/native/tile_reduce.h +394 -46
  165. warp/native/tile_scan.h +4 -4
  166. warp/native/vec.h +469 -53
  167. warp/native/version.h +23 -0
  168. warp/native/volume.cpp +1 -1
  169. warp/native/volume.cu +1 -0
  170. warp/native/volume.h +1 -1
  171. warp/native/volume_builder.cu +2 -0
  172. warp/native/warp.cpp +60 -32
  173. warp/native/warp.cu +313 -201
  174. warp/native/warp.h +14 -11
  175. warp/optim/__init__.py +6 -3
  176. warp/optim/adam.py +6 -145
  177. warp/optim/linear.py +14 -1585
  178. warp/optim/sgd.py +6 -94
  179. warp/paddle.py +6 -388
  180. warp/render/__init__.py +8 -4
  181. warp/render/imgui_manager.py +7 -267
  182. warp/render/render_opengl.py +6 -3616
  183. warp/render/render_usd.py +6 -918
  184. warp/render/utils.py +6 -142
  185. warp/sparse.py +37 -2563
  186. warp/tape.py +6 -1188
  187. warp/tests/__main__.py +1 -1
  188. warp/tests/cuda/test_async.py +4 -4
  189. warp/tests/cuda/test_conditional_captures.py +1 -1
  190. warp/tests/cuda/test_multigpu.py +1 -1
  191. warp/tests/cuda/test_streams.py +58 -1
  192. warp/tests/geometry/test_bvh.py +157 -22
  193. warp/tests/geometry/test_hash_grid.py +38 -0
  194. warp/tests/geometry/test_marching_cubes.py +0 -1
  195. warp/tests/geometry/test_mesh.py +5 -3
  196. warp/tests/geometry/test_mesh_query_aabb.py +5 -12
  197. warp/tests/geometry/test_mesh_query_point.py +5 -2
  198. warp/tests/geometry/test_mesh_query_ray.py +15 -3
  199. warp/tests/geometry/test_volume_write.py +5 -5
  200. warp/tests/interop/test_dlpack.py +14 -14
  201. warp/tests/interop/test_jax.py +1382 -79
  202. warp/tests/interop/test_paddle.py +1 -1
  203. warp/tests/test_adam.py +0 -1
  204. warp/tests/test_arithmetic.py +9 -9
  205. warp/tests/test_array.py +529 -100
  206. warp/tests/test_array_reduce.py +3 -3
  207. warp/tests/test_atomic.py +12 -8
  208. warp/tests/test_atomic_bitwise.py +209 -0
  209. warp/tests/test_atomic_cas.py +4 -4
  210. warp/tests/test_bool.py +2 -2
  211. warp/tests/test_builtins_resolution.py +5 -571
  212. warp/tests/test_codegen.py +34 -15
  213. warp/tests/test_conditional.py +1 -1
  214. warp/tests/test_context.py +6 -6
  215. warp/tests/test_copy.py +242 -161
  216. warp/tests/test_ctypes.py +3 -3
  217. warp/tests/test_devices.py +24 -2
  218. warp/tests/test_examples.py +16 -84
  219. warp/tests/test_fabricarray.py +35 -35
  220. warp/tests/test_fast_math.py +0 -2
  221. warp/tests/test_fem.py +60 -14
  222. warp/tests/test_fixedarray.py +3 -3
  223. warp/tests/test_func.py +8 -5
  224. warp/tests/test_generics.py +1 -1
  225. warp/tests/test_indexedarray.py +24 -24
  226. warp/tests/test_intersect.py +39 -9
  227. warp/tests/test_large.py +1 -1
  228. warp/tests/test_lerp.py +3 -1
  229. warp/tests/test_linear_solvers.py +1 -1
  230. warp/tests/test_map.py +49 -4
  231. warp/tests/test_mat.py +52 -62
  232. warp/tests/test_mat_constructors.py +4 -5
  233. warp/tests/test_mat_lite.py +1 -1
  234. warp/tests/test_mat_scalar_ops.py +121 -121
  235. warp/tests/test_math.py +34 -0
  236. warp/tests/test_module_aot.py +4 -4
  237. warp/tests/test_modules_lite.py +28 -2
  238. warp/tests/test_print.py +11 -11
  239. warp/tests/test_quat.py +93 -58
  240. warp/tests/test_runlength_encode.py +1 -1
  241. warp/tests/test_scalar_ops.py +38 -10
  242. warp/tests/test_smoothstep.py +1 -1
  243. warp/tests/test_sparse.py +126 -15
  244. warp/tests/test_spatial.py +105 -87
  245. warp/tests/test_special_values.py +6 -6
  246. warp/tests/test_static.py +7 -7
  247. warp/tests/test_struct.py +13 -2
  248. warp/tests/test_triangle_closest_point.py +48 -1
  249. warp/tests/test_tuple.py +96 -0
  250. warp/tests/test_types.py +82 -9
  251. warp/tests/test_utils.py +52 -52
  252. warp/tests/test_vec.py +29 -29
  253. warp/tests/test_vec_constructors.py +5 -5
  254. warp/tests/test_vec_scalar_ops.py +97 -97
  255. warp/tests/test_version.py +75 -0
  256. warp/tests/tile/test_tile.py +239 -0
  257. warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
  258. warp/tests/tile/test_tile_cholesky.py +7 -4
  259. warp/tests/tile/test_tile_load.py +26 -2
  260. warp/tests/tile/test_tile_mathdx.py +3 -3
  261. warp/tests/tile/test_tile_matmul.py +1 -1
  262. warp/tests/tile/test_tile_mlp.py +2 -4
  263. warp/tests/tile/test_tile_reduce.py +214 -13
  264. warp/tests/unittest_suites.py +6 -14
  265. warp/tests/unittest_utils.py +10 -9
  266. warp/tests/walkthrough_debug.py +3 -1
  267. warp/torch.py +6 -373
  268. warp/types.py +29 -5750
  269. warp/utils.py +10 -1659
  270. {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/METADATA +47 -103
  271. warp_lang-1.10.0rc2.dist-info/RECORD +468 -0
  272. warp_lang-1.10.0rc2.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
  273. warp_lang-1.10.0rc2.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
  274. warp_lang-1.10.0rc2.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
  275. warp_lang-1.10.0rc2.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
  276. warp_lang-1.10.0rc2.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
  277. warp_lang-1.10.0rc2.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
  278. warp_lang-1.10.0rc2.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
  279. warp_lang-1.10.0rc2.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
  280. warp_lang-1.10.0rc2.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
  281. warp_lang-1.10.0rc2.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
  282. warp_lang-1.10.0rc2.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
  283. warp_lang-1.10.0rc2.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
  284. warp_lang-1.10.0rc2.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
  285. warp_lang-1.10.0rc2.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
  286. warp_lang-1.10.0rc2.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
  287. warp/examples/assets/cartpole.urdf +0 -110
  288. warp/examples/assets/crazyflie.usd +0 -0
  289. warp/examples/assets/nv_ant.xml +0 -92
  290. warp/examples/assets/nv_humanoid.xml +0 -183
  291. warp/examples/assets/quadruped.urdf +0 -268
  292. warp/examples/optim/example_bounce.py +0 -266
  293. warp/examples/optim/example_cloth_throw.py +0 -228
  294. warp/examples/optim/example_drone.py +0 -870
  295. warp/examples/optim/example_inverse_kinematics.py +0 -182
  296. warp/examples/optim/example_inverse_kinematics_torch.py +0 -191
  297. warp/examples/optim/example_softbody_properties.py +0 -400
  298. warp/examples/optim/example_spring_cage.py +0 -245
  299. warp/examples/optim/example_trajectory.py +0 -227
  300. warp/examples/sim/example_cartpole.py +0 -143
  301. warp/examples/sim/example_cloth.py +0 -225
  302. warp/examples/sim/example_cloth_self_contact.py +0 -316
  303. warp/examples/sim/example_granular.py +0 -130
  304. warp/examples/sim/example_granular_collision_sdf.py +0 -202
  305. warp/examples/sim/example_jacobian_ik.py +0 -244
  306. warp/examples/sim/example_particle_chain.py +0 -124
  307. warp/examples/sim/example_quadruped.py +0 -203
  308. warp/examples/sim/example_rigid_chain.py +0 -203
  309. warp/examples/sim/example_rigid_contact.py +0 -195
  310. warp/examples/sim/example_rigid_force.py +0 -133
  311. warp/examples/sim/example_rigid_gyroscopic.py +0 -115
  312. warp/examples/sim/example_rigid_soft_contact.py +0 -140
  313. warp/examples/sim/example_soft_body.py +0 -196
  314. warp/examples/tile/example_tile_walker.py +0 -327
  315. warp/sim/__init__.py +0 -74
  316. warp/sim/articulation.py +0 -793
  317. warp/sim/collide.py +0 -2570
  318. warp/sim/graph_coloring.py +0 -307
  319. warp/sim/import_mjcf.py +0 -791
  320. warp/sim/import_snu.py +0 -227
  321. warp/sim/import_urdf.py +0 -579
  322. warp/sim/import_usd.py +0 -898
  323. warp/sim/inertia.py +0 -357
  324. warp/sim/integrator.py +0 -245
  325. warp/sim/integrator_euler.py +0 -2000
  326. warp/sim/integrator_featherstone.py +0 -2101
  327. warp/sim/integrator_vbd.py +0 -2487
  328. warp/sim/integrator_xpbd.py +0 -3295
  329. warp/sim/model.py +0 -4821
  330. warp/sim/particles.py +0 -121
  331. warp/sim/render.py +0 -431
  332. warp/sim/utils.py +0 -431
  333. warp/tests/sim/disabled_kinematics.py +0 -244
  334. warp/tests/sim/test_cloth.py +0 -863
  335. warp/tests/sim/test_collision.py +0 -743
  336. warp/tests/sim/test_coloring.py +0 -347
  337. warp/tests/sim/test_inertia.py +0 -161
  338. warp/tests/sim/test_model.py +0 -226
  339. warp/tests/sim/test_sim_grad.py +0 -287
  340. warp/tests/sim/test_sim_grad_bounce_linear.py +0 -212
  341. warp/tests/sim/test_sim_kinematics.py +0 -98
  342. warp/thirdparty/__init__.py +0 -0
  343. warp_lang-1.9.0.dist-info/RECORD +0 -456
  344. /warp/{fem → _src/fem}/quadrature/__init__.py +0 -0
  345. /warp/{tests/sim → _src/thirdparty}/__init__.py +0 -0
  346. /warp/{thirdparty → _src/thirdparty}/appdirs.py +0 -0
  347. /warp/{thirdparty → _src/thirdparty}/dlpack.py +0 -0
  348. {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/WHEEL +0 -0
  349. {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/licenses/LICENSE.md +0 -0
  350. {warp_lang-1.9.0.dist-info → warp_lang-1.10.0rc2.dist-info}/top_level.txt +0 -0
warp/build_dll.py CHANGED
@@ -1,4 +1,4 @@
1
- # SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
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,477 +13,12 @@
13
13
  # See the License for the specific language governing permissions and
14
14
  # limitations under the License.
15
15
 
16
- from __future__ import annotations
16
+ # TODO: Remove after cleaning up the public API.
17
17
 
18
- import os
19
- import platform
20
- import subprocess
21
- import sys
18
+ from warp._src import build_dll as _build_dll
22
19
 
23
- from warp.utils import ScopedTimer
24
20
 
25
- verbose_cmd = True # print command lines before executing them
21
+ def __getattr__(name):
22
+ from warp._src.utils import get_deprecated_api
26
23
 
27
- MIN_CTK_VERSION = (12, 0)
28
-
29
-
30
- def machine_architecture() -> str:
31
- """Return a canonical machine architecture string.
32
- - "x86_64" for x86-64, aka. AMD64, aka. x64
33
- - "aarch64" for AArch64, aka. ARM64
34
- """
35
- machine = platform.machine()
36
- if machine == "x86_64" or machine == "AMD64":
37
- return "x86_64"
38
- if machine == "aarch64" or machine == "arm64":
39
- return "aarch64"
40
- raise RuntimeError(f"Unrecognized machine architecture {machine}")
41
-
42
-
43
- def run_cmd(cmd):
44
- if verbose_cmd:
45
- print(cmd)
46
-
47
- try:
48
- return subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
49
- except subprocess.CalledProcessError as e:
50
- print("Command failed with exit code:", e.returncode)
51
- print("Command output was:")
52
- print(e.output.decode())
53
- raise e
54
-
55
-
56
- # cut-down version of vcvars64.bat that allows using
57
- # custom toolchain locations, returns the compiler program path
58
- def set_msvc_env(msvc_path, sdk_path):
59
- if "INCLUDE" not in os.environ:
60
- os.environ["INCLUDE"] = ""
61
-
62
- if "LIB" not in os.environ:
63
- os.environ["LIB"] = ""
64
-
65
- msvc_path = os.path.abspath(msvc_path)
66
- sdk_path = os.path.abspath(sdk_path)
67
-
68
- os.environ["INCLUDE"] += os.pathsep + os.path.join(msvc_path, "include")
69
- os.environ["INCLUDE"] += os.pathsep + os.path.join(sdk_path, "include/winrt")
70
- os.environ["INCLUDE"] += os.pathsep + os.path.join(sdk_path, "include/um")
71
- os.environ["INCLUDE"] += os.pathsep + os.path.join(sdk_path, "include/ucrt")
72
- os.environ["INCLUDE"] += os.pathsep + os.path.join(sdk_path, "include/shared")
73
-
74
- os.environ["LIB"] += os.pathsep + os.path.join(msvc_path, "lib/x64")
75
- os.environ["LIB"] += os.pathsep + os.path.join(sdk_path, "lib/ucrt/x64")
76
- os.environ["LIB"] += os.pathsep + os.path.join(sdk_path, "lib/um/x64")
77
-
78
- os.environ["PATH"] += os.pathsep + os.path.join(msvc_path, "bin/HostX64/x64")
79
- os.environ["PATH"] += os.pathsep + os.path.join(sdk_path, "bin/x64")
80
-
81
- return os.path.join(msvc_path, "bin", "HostX64", "x64", "cl.exe")
82
-
83
-
84
- def find_host_compiler():
85
- if os.name == "nt":
86
- # try and find an installed host compiler (msvc)
87
- # runs vcvars and copies back the build environment
88
-
89
- vswhere_path = r"%ProgramFiles(x86)%/Microsoft Visual Studio/Installer/vswhere.exe"
90
- vswhere_path = os.path.expandvars(vswhere_path)
91
- if not os.path.exists(vswhere_path):
92
- return ""
93
-
94
- vs_path = run_cmd(f'"{vswhere_path}" -latest -property installationPath').decode().rstrip()
95
- vsvars_path = os.path.join(vs_path, "VC\\Auxiliary\\Build\\vcvars64.bat")
96
-
97
- output = run_cmd(f'"{vsvars_path}" && set').decode()
98
-
99
- for line in output.splitlines():
100
- pair = line.split("=", 1)
101
- if len(pair) >= 2:
102
- os.environ[pair[0]] = pair[1]
103
-
104
- cl_path = run_cmd("where cl.exe").decode("utf-8").rstrip()
105
- cl_version = os.environ["VCToolsVersion"].split(".")
106
-
107
- # ensure at least VS2019 version, see list of MSVC versions here https://en.wikipedia.org/wiki/Microsoft_Visual_C%2B%2B
108
- cl_required_major = 14
109
- cl_required_minor = 29
110
-
111
- if int(cl_version[0]) < cl_required_major or (
112
- (int(cl_version[0]) == cl_required_major) and (int(cl_version[1]) < cl_required_minor)
113
- ):
114
- print(
115
- f"Warp: MSVC found but compiler version too old, found {cl_version[0]}.{cl_version[1]}, but must be {cl_required_major}.{cl_required_minor} or higher, kernel host compilation will be disabled."
116
- )
117
- return ""
118
-
119
- return cl_path
120
-
121
- else:
122
- # try and find g++
123
- return run_cmd("which g++").decode()
124
-
125
-
126
- def get_cuda_toolkit_version(cuda_home) -> tuple[int, int]:
127
- try:
128
- # the toolkit version can be obtained by running "nvcc --version"
129
- nvcc_path = os.path.join(cuda_home, "bin", "nvcc")
130
- nvcc_version_output = subprocess.check_output([nvcc_path, "--version"]).decode("utf-8")
131
- # search for release substring (e.g., "release 11.5")
132
- import re
133
-
134
- m = re.search(r"release (\d+)\.(\d+)", nvcc_version_output)
135
- if m is not None:
136
- major, minor = map(int, m.groups())
137
- return (major, minor)
138
- else:
139
- raise Exception("Failed to parse NVCC output")
140
-
141
- except Exception as e:
142
- print(f"Warning: Failed to determine CUDA Toolkit version: {e}")
143
- return MIN_CTK_VERSION
144
-
145
-
146
- def quote(path):
147
- return '"' + path + '"'
148
-
149
-
150
- def add_llvm_bin_to_path(args):
151
- """Add the LLVM bin directory to the PATH environment variable if it's set.
152
-
153
- Args:
154
- args: The argument namespace containing llvm_path.
155
-
156
- Returns:
157
- ``True`` if the PATH was updated, ``False`` otherwise.
158
- """
159
- if not hasattr(args, "llvm_path") or not args.llvm_path:
160
- return False
161
-
162
- # Construct the bin directory path
163
- llvm_bin_path = os.path.join(args.llvm_path, "bin")
164
-
165
- # Check if the directory exists
166
- if not os.path.isdir(llvm_bin_path):
167
- print(f"Warning: LLVM bin directory not found at {llvm_bin_path}")
168
- return False
169
-
170
- # Add to PATH environment variable
171
- os.environ["PATH"] = llvm_bin_path + os.pathsep + os.environ.get("PATH", "")
172
-
173
- print(f"Added {llvm_bin_path} to PATH")
174
- return True
175
-
176
-
177
- def build_dll_for_arch(args, dll_path, cpp_paths, cu_path, arch, libs: list[str] | None = None, mode=None):
178
- mode = args.mode if (mode is None) else mode
179
- cuda_home = args.cuda_path
180
- cuda_cmd = None
181
-
182
- # Add LLVM bin directory to PATH
183
- add_llvm_bin_to_path(args)
184
-
185
- if args.quick or cu_path is None:
186
- cuda_compat_enabled = "WP_ENABLE_CUDA_COMPATIBILITY=0"
187
- else:
188
- cuda_compat_enabled = "WP_ENABLE_CUDA_COMPATIBILITY=1"
189
-
190
- if libs is None:
191
- libs = []
192
-
193
- import pathlib
194
-
195
- warp_home_path = pathlib.Path(__file__).parent
196
- warp_home = warp_home_path.resolve()
197
-
198
- if args.verbose:
199
- print(f"Building {dll_path}")
200
-
201
- native_dir = os.path.join(warp_home, "native")
202
-
203
- if cu_path:
204
- # check CUDA Toolkit version
205
- ctk_version = get_cuda_toolkit_version(cuda_home)
206
- if ctk_version < MIN_CTK_VERSION:
207
- raise Exception(
208
- f"CUDA Toolkit version {MIN_CTK_VERSION[0]}.{MIN_CTK_VERSION[1]}+ is required (found {ctk_version[0]}.{ctk_version[1]} in {cuda_home})"
209
- )
210
-
211
- # NVCC gencode options
212
- gencode_opts = []
213
-
214
- # Clang architecture flags
215
- clang_arch_flags = []
216
-
217
- if args.quick:
218
- # minimum supported architectures (PTX)
219
- if ctk_version >= (13, 0):
220
- gencode_opts += ["-gencode=arch=compute_75,code=compute_75"]
221
- clang_arch_flags += ["--cuda-gpu-arch=sm_75"]
222
- else:
223
- gencode_opts += ["-gencode=arch=compute_52,code=compute_52", "-gencode=arch=compute_75,code=compute_75"]
224
- clang_arch_flags += ["--cuda-gpu-arch=sm_52", "--cuda-gpu-arch=sm_75"]
225
- else:
226
- # generate code for all supported architectures
227
- gencode_opts += [
228
- # SASS for supported desktop/datacenter architectures
229
- "-gencode=arch=compute_75,code=sm_75", # Turing
230
- "-gencode=arch=compute_75,code=compute_75", # Turing (PTX)
231
- "-gencode=arch=compute_80,code=sm_80", # Ampere
232
- "-gencode=arch=compute_86,code=sm_86",
233
- "-gencode=arch=compute_89,code=sm_89", # Ada
234
- "-gencode=arch=compute_90,code=sm_90", # Hopper
235
- ]
236
-
237
- clang_arch_flags += [
238
- # SASS for supported desktop/datacenter architectures
239
- "--cuda-gpu-arch=sm_75", # Turing
240
- "--cuda-gpu-arch=sm_80", # Ampere
241
- "--cuda-gpu-arch=sm_86",
242
- "--cuda-gpu-arch=sm_89", # Ada
243
- "--cuda-gpu-arch=sm_90", # Hopper
244
- ]
245
-
246
- if arch == "aarch64" and sys.platform == "linux":
247
- # SASS for supported mobile architectures (e.g. Tegra/Jetson)
248
- gencode_opts += ["-gencode=arch=compute_87,code=sm_87"] # Orin
249
- clang_arch_flags += ["--cuda-gpu-arch=sm_87"]
250
-
251
- if ctk_version >= (13, 0):
252
- gencode_opts += ["-gencode=arch=compute_110,code=sm_110"] # Thor
253
- clang_arch_flags += ["--cuda-gpu-arch=sm_110"]
254
- else:
255
- gencode_opts += [
256
- "-gencode=arch=compute_53,code=sm_53", # X1
257
- "-gencode=arch=compute_62,code=sm_62", # X2
258
- "-gencode=arch=compute_72,code=sm_72", # Xavier
259
- ]
260
- clang_arch_flags += [
261
- "--cuda-gpu-arch=sm_53",
262
- "--cuda-gpu-arch=sm_62",
263
- "--cuda-gpu-arch=sm_72",
264
- ]
265
-
266
- if ctk_version >= (12, 8):
267
- gencode_opts += ["-gencode=arch=compute_101,code=sm_101"] # Thor (CUDA 12 numbering)
268
- clang_arch_flags += ["--cuda-gpu-arch=sm_101"]
269
-
270
- if ctk_version >= (12, 8):
271
- # Support for Blackwell is available with CUDA Toolkit 12.8+
272
- gencode_opts += [
273
- "-gencode=arch=compute_100,code=sm_100", # Blackwell
274
- "-gencode=arch=compute_120,code=sm_120", # Blackwell
275
- "-gencode=arch=compute_120,code=compute_120", # PTX for future hardware
276
- ]
277
-
278
- clang_arch_flags += [
279
- "--cuda-gpu-arch=sm_100", # Blackwell
280
- "--cuda-gpu-arch=sm_120", # Blackwell
281
- ]
282
- else:
283
- gencode_opts += ["-gencode=arch=compute_90,code=compute_90"] # PTX for future hardware
284
-
285
- nvcc_opts = [
286
- *gencode_opts,
287
- "-t0", # multithreaded compilation
288
- "--extended-lambda",
289
- ]
290
-
291
- # Clang options
292
- clang_opts = [
293
- *clang_arch_flags,
294
- "-std=c++17",
295
- "-xcuda",
296
- f'--cuda-path="{cuda_home}"',
297
- ]
298
-
299
- if args.compile_time_trace:
300
- if ctk_version >= (12, 8):
301
- nvcc_opts.append("--fdevice-time-trace=build_lib_compile-time-trace")
302
- else:
303
- print("Warp warning: CUDA version is less than 12.8, compile_time_trace is not supported")
304
-
305
- if args.fast_math:
306
- nvcc_opts.append("--use_fast_math")
307
-
308
- # is the library being built with CUDA enabled?
309
- cuda_enabled = "WP_ENABLE_CUDA=1" if (cu_path is not None) else "WP_ENABLE_CUDA=0"
310
-
311
- if args.libmathdx_path:
312
- libmathdx_includes = f' -I"{args.libmathdx_path}/include"'
313
- mathdx_enabled = "WP_ENABLE_MATHDX=1"
314
- else:
315
- libmathdx_includes = ""
316
- mathdx_enabled = "WP_ENABLE_MATHDX=0"
317
-
318
- if os.name == "nt":
319
- if args.host_compiler:
320
- host_linker = os.path.join(os.path.dirname(args.host_compiler), "link.exe")
321
- else:
322
- raise RuntimeError("Warp build error: No host compiler was found")
323
-
324
- cpp_includes = f' /I"{warp_home_path.parent}/external/llvm-project/out/install/{mode}-{arch}/include"'
325
- cpp_includes += f' /I"{warp_home_path.parent}/_build/host-deps/llvm-project/release-{arch}/include"'
326
- cuda_includes = f' /I"{cuda_home}/include"' if cu_path else ""
327
- includes = cpp_includes + cuda_includes
328
-
329
- # nvrtc_static.lib is built with /MT and _ITERATOR_DEBUG_LEVEL=0 so if we link it in we must match these options
330
- if cu_path or mode != "debug":
331
- runtime = "/MT"
332
- iter_dbg = "_ITERATOR_DEBUG_LEVEL=0"
333
- debug = "NDEBUG"
334
- else:
335
- runtime = "/MTd"
336
- iter_dbg = "_ITERATOR_DEBUG_LEVEL=2"
337
- debug = "_DEBUG"
338
-
339
- cpp_flags = f'/nologo /std:c++17 /GR- {runtime} /D "{debug}" /D "{cuda_enabled}" /D "{mathdx_enabled}" /D "{cuda_compat_enabled}" /D "{iter_dbg}" /I"{native_dir}" {includes} '
340
-
341
- if args.mode == "debug":
342
- cpp_flags += "/Zi /Od /D WP_ENABLE_DEBUG=1"
343
- linkopts = ["/DLL", "/DEBUG"]
344
- elif args.mode == "release":
345
- cpp_flags += "/Ox /D WP_ENABLE_DEBUG=0"
346
- linkopts = ["/DLL"]
347
- else:
348
- raise RuntimeError(f"Unrecognized build configuration (debug, release), got: {args.mode}")
349
-
350
- if args.verify_fp:
351
- cpp_flags += ' /D "WP_VERIFY_FP"'
352
-
353
- if args.fast_math:
354
- cpp_flags += " /fp:fast"
355
-
356
- with ScopedTimer("build", active=args.verbose):
357
- for cpp_path in cpp_paths:
358
- cpp_out = cpp_path + ".obj"
359
- linkopts.append(quote(cpp_out))
360
-
361
- cpp_cmd = f'"{args.host_compiler}" {cpp_flags} -c "{cpp_path}" /Fo"{cpp_out}"'
362
- run_cmd(cpp_cmd)
363
-
364
- if cu_path:
365
- cu_out = cu_path + ".o"
366
-
367
- if mode == "debug":
368
- cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 --compiler-options=/MT,/Zi,/Od -g -G -O0 -DNDEBUG -D_ITERATOR_DEBUG_LEVEL=0 -I"{native_dir}" -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
369
-
370
- elif mode == "release":
371
- cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 {" ".join(nvcc_opts)} -I"{native_dir}" -DNDEBUG -DWP_ENABLE_CUDA=1 -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
372
-
373
- with ScopedTimer("build_cuda", active=args.verbose):
374
- run_cmd(cuda_cmd)
375
- linkopts.append(quote(cu_out))
376
- linkopts.append(
377
- f'cudart_static.lib nvrtc_static.lib nvrtc-builtins_static.lib nvptxcompiler_static.lib ws2_32.lib user32.lib /LIBPATH:"{cuda_home}/lib/x64"'
378
- )
379
-
380
- if args.libmathdx_path:
381
- linkopts.append(f'nvJitLink_static.lib /LIBPATH:"{args.libmathdx_path}/lib/x64" mathdx_static.lib')
382
-
383
- with ScopedTimer("link", active=args.verbose):
384
- link_cmd = f'"{host_linker}" {" ".join(linkopts + libs)} /out:"{dll_path}"'
385
- run_cmd(link_cmd)
386
-
387
- else:
388
- # Unix compilation
389
- cuda_compiler = "clang++" if getattr(args, "clang_build_toolchain", False) else "nvcc"
390
- cpp_compiler = "clang++" if getattr(args, "clang_build_toolchain", False) else "g++"
391
-
392
- cpp_includes = f' -I"{warp_home_path.parent}/external/llvm-project/out/install/{mode}-{arch}/include"'
393
- cpp_includes += f' -I"{warp_home_path.parent}/_build/host-deps/llvm-project/release-{arch}/include"'
394
- cuda_includes = f' -I"{cuda_home}/include"' if cu_path else ""
395
- includes = cpp_includes + cuda_includes
396
-
397
- if sys.platform == "darwin":
398
- version = f"--target={arch}-apple-macos11"
399
- else:
400
- if cpp_compiler == "g++":
401
- version = "-fabi-version=13" # GCC 8.2+
402
- else:
403
- version = ""
404
-
405
- cpp_flags = f'-Werror -Wuninitialized {version} --std=c++17 -fno-rtti -D{cuda_enabled} -D{mathdx_enabled} -D{cuda_compat_enabled} -fPIC -fvisibility=hidden -D_GLIBCXX_USE_CXX11_ABI=0 -I"{native_dir}" {includes} '
406
-
407
- if mode == "debug":
408
- cpp_flags += "-O0 -g -D_DEBUG -DWP_ENABLE_DEBUG=1 -fkeep-inline-functions"
409
-
410
- if mode == "release":
411
- cpp_flags += "-O3 -DNDEBUG -DWP_ENABLE_DEBUG=0"
412
-
413
- if args.verify_fp:
414
- cpp_flags += " -DWP_VERIFY_FP"
415
-
416
- if args.fast_math:
417
- cpp_flags += " -ffast-math"
418
-
419
- ld_inputs = []
420
-
421
- with ScopedTimer("build", active=args.verbose):
422
- for cpp_path in cpp_paths:
423
- cpp_out = cpp_path + ".o"
424
- ld_inputs.append(quote(cpp_out))
425
-
426
- build_cmd = f'{cpp_compiler} {cpp_flags} -c "{cpp_path}" -o "{cpp_out}"'
427
- run_cmd(build_cmd)
428
-
429
- if cu_path:
430
- cu_out = cu_path + ".o"
431
-
432
- if cuda_compiler == "nvcc":
433
- if mode == "debug":
434
- cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -g -G -O0 --compiler-options -fPIC,-fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -line-info {" ".join(nvcc_opts)} -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
435
- elif mode == "release":
436
- cuda_cmd = f'"{cuda_home}/bin/nvcc" --std=c++17 -O3 --compiler-options -fPIC,-fvisibility=hidden {" ".join(nvcc_opts)} -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
437
- else:
438
- # Use Clang compiler
439
- if mode == "debug":
440
- cuda_cmd = f'clang++ -Werror -Wuninitialized -Wno-unknown-cuda-version {" ".join(clang_opts)} -g -O0 -fPIC -fvisibility=hidden -D_DEBUG -D_ITERATOR_DEBUG_LEVEL=0 -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
441
- elif mode == "release":
442
- cuda_cmd = f'clang++ -Werror -Wuninitialized -Wno-unknown-cuda-version {" ".join(clang_opts)} -O3 -fPIC -fvisibility=hidden -DNDEBUG -DWP_ENABLE_CUDA=1 -I"{native_dir}" -D{mathdx_enabled} {libmathdx_includes} -o "{cu_out}" -c "{cu_path}"'
443
-
444
- with ScopedTimer("build_cuda", active=args.verbose):
445
- run_cmd(cuda_cmd)
446
-
447
- ld_inputs.append(quote(cu_out))
448
- ld_inputs.append(
449
- f'-L"{cuda_home}/lib64" -lcudart_static -lnvrtc_static -lnvrtc-builtins_static -lnvptxcompiler_static -lpthread -ldl -lrt'
450
- )
451
-
452
- if args.libmathdx_path:
453
- ld_inputs.append(f"-lnvJitLink_static -L{args.libmathdx_path}/lib -lmathdx_static")
454
-
455
- if sys.platform == "darwin":
456
- opt_no_undefined = "-Wl,-undefined,error"
457
- opt_exclude_libs = ""
458
- else:
459
- opt_no_undefined = "-Wl,--no-undefined"
460
- opt_exclude_libs = "-Wl,--exclude-libs,ALL"
461
-
462
- with ScopedTimer("link", active=args.verbose):
463
- origin = "@loader_path" if (sys.platform == "darwin") else "$ORIGIN"
464
- link_cmd = f"{cpp_compiler} {version} -shared -Wl,-rpath,'{origin}' {opt_no_undefined} {opt_exclude_libs} -o '{dll_path}' {' '.join(ld_inputs + libs)}"
465
- run_cmd(link_cmd)
466
-
467
- # Strip symbols to reduce the binary size
468
- if mode == "release":
469
- if sys.platform == "darwin":
470
- run_cmd(f"strip -x {dll_path}") # Strip all local symbols
471
- else: # Linux
472
- # Strip all symbols except for those needed to support debugging JIT-compiled code
473
- run_cmd(
474
- f"strip --strip-all --keep-symbol=__jit_debug_register_code --keep-symbol=__jit_debug_descriptor {dll_path}"
475
- )
476
-
477
-
478
- def build_dll(args, dll_path, cpp_paths, cu_path, libs=None):
479
- if sys.platform == "darwin":
480
- # create a universal binary by combining x86-64 and AArch64 builds
481
- build_dll_for_arch(args, dll_path + "-x86_64", cpp_paths, cu_path, "x86_64", libs)
482
- build_dll_for_arch(args, dll_path + "-aarch64", cpp_paths, cu_path, "aarch64", libs)
483
-
484
- run_cmd(f"lipo -create -output {dll_path} {dll_path}-x86_64 {dll_path}-aarch64")
485
- os.remove(f"{dll_path}-x86_64")
486
- os.remove(f"{dll_path}-aarch64")
487
-
488
- else:
489
- build_dll_for_arch(args, dll_path, cpp_paths, cu_path, machine_architecture(), libs)
24
+ return get_deprecated_api(_build_dll, "wp", name)