warp-lang 1.10.0__py3-none-macosx_11_0_arm64.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 (468) hide show
  1. warp/__init__.py +334 -0
  2. warp/__init__.pyi +5856 -0
  3. warp/_src/__init__.py +14 -0
  4. warp/_src/autograd.py +1077 -0
  5. warp/_src/build.py +620 -0
  6. warp/_src/build_dll.py +642 -0
  7. warp/_src/builtins.py +10555 -0
  8. warp/_src/codegen.py +4361 -0
  9. warp/_src/config.py +178 -0
  10. warp/_src/constants.py +59 -0
  11. warp/_src/context.py +8352 -0
  12. warp/_src/dlpack.py +464 -0
  13. warp/_src/fabric.py +362 -0
  14. warp/_src/fem/__init__.py +14 -0
  15. warp/_src/fem/adaptivity.py +510 -0
  16. warp/_src/fem/cache.py +689 -0
  17. warp/_src/fem/dirichlet.py +190 -0
  18. warp/_src/fem/domain.py +553 -0
  19. warp/_src/fem/field/__init__.py +131 -0
  20. warp/_src/fem/field/field.py +703 -0
  21. warp/_src/fem/field/nodal_field.py +403 -0
  22. warp/_src/fem/field/restriction.py +39 -0
  23. warp/_src/fem/field/virtual.py +1021 -0
  24. warp/_src/fem/geometry/__init__.py +32 -0
  25. warp/_src/fem/geometry/adaptive_nanogrid.py +782 -0
  26. warp/_src/fem/geometry/closest_point.py +99 -0
  27. warp/_src/fem/geometry/deformed_geometry.py +277 -0
  28. warp/_src/fem/geometry/element.py +854 -0
  29. warp/_src/fem/geometry/geometry.py +693 -0
  30. warp/_src/fem/geometry/grid_2d.py +478 -0
  31. warp/_src/fem/geometry/grid_3d.py +539 -0
  32. warp/_src/fem/geometry/hexmesh.py +956 -0
  33. warp/_src/fem/geometry/nanogrid.py +660 -0
  34. warp/_src/fem/geometry/partition.py +483 -0
  35. warp/_src/fem/geometry/quadmesh.py +597 -0
  36. warp/_src/fem/geometry/tetmesh.py +762 -0
  37. warp/_src/fem/geometry/trimesh.py +588 -0
  38. warp/_src/fem/integrate.py +2507 -0
  39. warp/_src/fem/linalg.py +385 -0
  40. warp/_src/fem/operator.py +398 -0
  41. warp/_src/fem/polynomial.py +231 -0
  42. warp/_src/fem/quadrature/__init__.py +17 -0
  43. warp/_src/fem/quadrature/pic_quadrature.py +318 -0
  44. warp/_src/fem/quadrature/quadrature.py +665 -0
  45. warp/_src/fem/space/__init__.py +248 -0
  46. warp/_src/fem/space/basis_function_space.py +499 -0
  47. warp/_src/fem/space/basis_space.py +681 -0
  48. warp/_src/fem/space/dof_mapper.py +253 -0
  49. warp/_src/fem/space/function_space.py +312 -0
  50. warp/_src/fem/space/grid_2d_function_space.py +179 -0
  51. warp/_src/fem/space/grid_3d_function_space.py +229 -0
  52. warp/_src/fem/space/hexmesh_function_space.py +255 -0
  53. warp/_src/fem/space/nanogrid_function_space.py +199 -0
  54. warp/_src/fem/space/partition.py +435 -0
  55. warp/_src/fem/space/quadmesh_function_space.py +222 -0
  56. warp/_src/fem/space/restriction.py +221 -0
  57. warp/_src/fem/space/shape/__init__.py +152 -0
  58. warp/_src/fem/space/shape/cube_shape_function.py +1107 -0
  59. warp/_src/fem/space/shape/shape_function.py +134 -0
  60. warp/_src/fem/space/shape/square_shape_function.py +928 -0
  61. warp/_src/fem/space/shape/tet_shape_function.py +829 -0
  62. warp/_src/fem/space/shape/triangle_shape_function.py +674 -0
  63. warp/_src/fem/space/tetmesh_function_space.py +270 -0
  64. warp/_src/fem/space/topology.py +461 -0
  65. warp/_src/fem/space/trimesh_function_space.py +193 -0
  66. warp/_src/fem/types.py +114 -0
  67. warp/_src/fem/utils.py +488 -0
  68. warp/_src/jax.py +188 -0
  69. warp/_src/jax_experimental/__init__.py +14 -0
  70. warp/_src/jax_experimental/custom_call.py +389 -0
  71. warp/_src/jax_experimental/ffi.py +1286 -0
  72. warp/_src/jax_experimental/xla_ffi.py +658 -0
  73. warp/_src/marching_cubes.py +710 -0
  74. warp/_src/math.py +416 -0
  75. warp/_src/optim/__init__.py +14 -0
  76. warp/_src/optim/adam.py +165 -0
  77. warp/_src/optim/linear.py +1608 -0
  78. warp/_src/optim/sgd.py +114 -0
  79. warp/_src/paddle.py +408 -0
  80. warp/_src/render/__init__.py +14 -0
  81. warp/_src/render/imgui_manager.py +291 -0
  82. warp/_src/render/render_opengl.py +3638 -0
  83. warp/_src/render/render_usd.py +939 -0
  84. warp/_src/render/utils.py +162 -0
  85. warp/_src/sparse.py +2718 -0
  86. warp/_src/tape.py +1208 -0
  87. warp/_src/thirdparty/__init__.py +0 -0
  88. warp/_src/thirdparty/appdirs.py +598 -0
  89. warp/_src/thirdparty/dlpack.py +145 -0
  90. warp/_src/thirdparty/unittest_parallel.py +676 -0
  91. warp/_src/torch.py +393 -0
  92. warp/_src/types.py +5888 -0
  93. warp/_src/utils.py +1695 -0
  94. warp/autograd.py +33 -0
  95. warp/bin/libwarp-clang.dylib +0 -0
  96. warp/bin/libwarp.dylib +0 -0
  97. warp/build.py +29 -0
  98. warp/build_dll.py +24 -0
  99. warp/codegen.py +24 -0
  100. warp/constants.py +24 -0
  101. warp/context.py +33 -0
  102. warp/dlpack.py +24 -0
  103. warp/examples/__init__.py +24 -0
  104. warp/examples/assets/bear.usd +0 -0
  105. warp/examples/assets/bunny.usd +0 -0
  106. warp/examples/assets/cube.usd +0 -0
  107. warp/examples/assets/nonuniform.usd +0 -0
  108. warp/examples/assets/nvidia_logo.png +0 -0
  109. warp/examples/assets/pixel.jpg +0 -0
  110. warp/examples/assets/rocks.nvdb +0 -0
  111. warp/examples/assets/rocks.usd +0 -0
  112. warp/examples/assets/sphere.usd +0 -0
  113. warp/examples/assets/square_cloth.usd +0 -0
  114. warp/examples/benchmarks/benchmark_api.py +389 -0
  115. warp/examples/benchmarks/benchmark_cloth.py +296 -0
  116. warp/examples/benchmarks/benchmark_cloth_cupy.py +96 -0
  117. warp/examples/benchmarks/benchmark_cloth_jax.py +105 -0
  118. warp/examples/benchmarks/benchmark_cloth_numba.py +161 -0
  119. warp/examples/benchmarks/benchmark_cloth_numpy.py +85 -0
  120. warp/examples/benchmarks/benchmark_cloth_paddle.py +94 -0
  121. warp/examples/benchmarks/benchmark_cloth_pytorch.py +94 -0
  122. warp/examples/benchmarks/benchmark_cloth_taichi.py +120 -0
  123. warp/examples/benchmarks/benchmark_cloth_warp.py +153 -0
  124. warp/examples/benchmarks/benchmark_gemm.py +164 -0
  125. warp/examples/benchmarks/benchmark_interop_paddle.py +166 -0
  126. warp/examples/benchmarks/benchmark_interop_torch.py +166 -0
  127. warp/examples/benchmarks/benchmark_launches.py +301 -0
  128. warp/examples/benchmarks/benchmark_tile_load_store.py +103 -0
  129. warp/examples/benchmarks/benchmark_tile_sort.py +155 -0
  130. warp/examples/browse.py +37 -0
  131. warp/examples/core/example_cupy.py +86 -0
  132. warp/examples/core/example_dem.py +241 -0
  133. warp/examples/core/example_fluid.py +299 -0
  134. warp/examples/core/example_graph_capture.py +150 -0
  135. warp/examples/core/example_marching_cubes.py +195 -0
  136. warp/examples/core/example_mesh.py +180 -0
  137. warp/examples/core/example_mesh_intersect.py +211 -0
  138. warp/examples/core/example_nvdb.py +182 -0
  139. warp/examples/core/example_raycast.py +111 -0
  140. warp/examples/core/example_raymarch.py +205 -0
  141. warp/examples/core/example_render_opengl.py +290 -0
  142. warp/examples/core/example_sample_mesh.py +300 -0
  143. warp/examples/core/example_sph.py +411 -0
  144. warp/examples/core/example_spin_lock.py +93 -0
  145. warp/examples/core/example_torch.py +211 -0
  146. warp/examples/core/example_wave.py +269 -0
  147. warp/examples/core/example_work_queue.py +118 -0
  148. warp/examples/distributed/example_jacobi_mpi.py +506 -0
  149. warp/examples/fem/example_adaptive_grid.py +286 -0
  150. warp/examples/fem/example_apic_fluid.py +469 -0
  151. warp/examples/fem/example_burgers.py +261 -0
  152. warp/examples/fem/example_convection_diffusion.py +181 -0
  153. warp/examples/fem/example_convection_diffusion_dg.py +225 -0
  154. warp/examples/fem/example_darcy_ls_optimization.py +489 -0
  155. warp/examples/fem/example_deformed_geometry.py +172 -0
  156. warp/examples/fem/example_diffusion.py +196 -0
  157. warp/examples/fem/example_diffusion_3d.py +225 -0
  158. warp/examples/fem/example_diffusion_mgpu.py +225 -0
  159. warp/examples/fem/example_distortion_energy.py +228 -0
  160. warp/examples/fem/example_elastic_shape_optimization.py +387 -0
  161. warp/examples/fem/example_magnetostatics.py +242 -0
  162. warp/examples/fem/example_mixed_elasticity.py +293 -0
  163. warp/examples/fem/example_navier_stokes.py +263 -0
  164. warp/examples/fem/example_nonconforming_contact.py +300 -0
  165. warp/examples/fem/example_stokes.py +213 -0
  166. warp/examples/fem/example_stokes_transfer.py +262 -0
  167. warp/examples/fem/example_streamlines.py +357 -0
  168. warp/examples/fem/utils.py +1047 -0
  169. warp/examples/interop/example_jax_callable.py +146 -0
  170. warp/examples/interop/example_jax_ffi_callback.py +132 -0
  171. warp/examples/interop/example_jax_kernel.py +232 -0
  172. warp/examples/optim/example_diffray.py +561 -0
  173. warp/examples/optim/example_fluid_checkpoint.py +497 -0
  174. warp/examples/tile/example_tile_block_cholesky.py +502 -0
  175. warp/examples/tile/example_tile_cholesky.py +88 -0
  176. warp/examples/tile/example_tile_convolution.py +66 -0
  177. warp/examples/tile/example_tile_fft.py +55 -0
  178. warp/examples/tile/example_tile_filtering.py +113 -0
  179. warp/examples/tile/example_tile_matmul.py +85 -0
  180. warp/examples/tile/example_tile_mcgp.py +191 -0
  181. warp/examples/tile/example_tile_mlp.py +385 -0
  182. warp/examples/tile/example_tile_nbody.py +199 -0
  183. warp/fabric.py +24 -0
  184. warp/fem/__init__.py +173 -0
  185. warp/fem/adaptivity.py +26 -0
  186. warp/fem/cache.py +30 -0
  187. warp/fem/dirichlet.py +24 -0
  188. warp/fem/field/__init__.py +24 -0
  189. warp/fem/field/field.py +26 -0
  190. warp/fem/geometry/__init__.py +21 -0
  191. warp/fem/geometry/closest_point.py +31 -0
  192. warp/fem/linalg.py +38 -0
  193. warp/fem/operator.py +32 -0
  194. warp/fem/polynomial.py +29 -0
  195. warp/fem/space/__init__.py +22 -0
  196. warp/fem/space/basis_space.py +24 -0
  197. warp/fem/space/shape/__init__.py +68 -0
  198. warp/fem/space/topology.py +24 -0
  199. warp/fem/types.py +24 -0
  200. warp/fem/utils.py +32 -0
  201. warp/jax.py +29 -0
  202. warp/jax_experimental/__init__.py +29 -0
  203. warp/jax_experimental/custom_call.py +29 -0
  204. warp/jax_experimental/ffi.py +39 -0
  205. warp/jax_experimental/xla_ffi.py +24 -0
  206. warp/marching_cubes.py +24 -0
  207. warp/math.py +37 -0
  208. warp/native/array.h +1687 -0
  209. warp/native/builtin.h +2327 -0
  210. warp/native/bvh.cpp +562 -0
  211. warp/native/bvh.cu +826 -0
  212. warp/native/bvh.h +555 -0
  213. warp/native/clang/clang.cpp +541 -0
  214. warp/native/coloring.cpp +622 -0
  215. warp/native/crt.cpp +51 -0
  216. warp/native/crt.h +568 -0
  217. warp/native/cuda_crt.h +1058 -0
  218. warp/native/cuda_util.cpp +677 -0
  219. warp/native/cuda_util.h +313 -0
  220. warp/native/error.cpp +77 -0
  221. warp/native/error.h +36 -0
  222. warp/native/exports.h +2023 -0
  223. warp/native/fabric.h +246 -0
  224. warp/native/hashgrid.cpp +311 -0
  225. warp/native/hashgrid.cu +89 -0
  226. warp/native/hashgrid.h +240 -0
  227. warp/native/initializer_array.h +41 -0
  228. warp/native/intersect.h +1253 -0
  229. warp/native/intersect_adj.h +375 -0
  230. warp/native/intersect_tri.h +348 -0
  231. warp/native/mat.h +5189 -0
  232. warp/native/mathdx.cpp +93 -0
  233. warp/native/matnn.h +221 -0
  234. warp/native/mesh.cpp +266 -0
  235. warp/native/mesh.cu +406 -0
  236. warp/native/mesh.h +2097 -0
  237. warp/native/nanovdb/GridHandle.h +533 -0
  238. warp/native/nanovdb/HostBuffer.h +591 -0
  239. warp/native/nanovdb/NanoVDB.h +6246 -0
  240. warp/native/nanovdb/NodeManager.h +323 -0
  241. warp/native/nanovdb/PNanoVDB.h +3390 -0
  242. warp/native/noise.h +859 -0
  243. warp/native/quat.h +1664 -0
  244. warp/native/rand.h +342 -0
  245. warp/native/range.h +145 -0
  246. warp/native/reduce.cpp +174 -0
  247. warp/native/reduce.cu +363 -0
  248. warp/native/runlength_encode.cpp +79 -0
  249. warp/native/runlength_encode.cu +61 -0
  250. warp/native/scan.cpp +47 -0
  251. warp/native/scan.cu +55 -0
  252. warp/native/scan.h +23 -0
  253. warp/native/solid_angle.h +466 -0
  254. warp/native/sort.cpp +251 -0
  255. warp/native/sort.cu +286 -0
  256. warp/native/sort.h +35 -0
  257. warp/native/sparse.cpp +241 -0
  258. warp/native/sparse.cu +435 -0
  259. warp/native/spatial.h +1306 -0
  260. warp/native/svd.h +727 -0
  261. warp/native/temp_buffer.h +46 -0
  262. warp/native/tile.h +4124 -0
  263. warp/native/tile_radix_sort.h +1112 -0
  264. warp/native/tile_reduce.h +838 -0
  265. warp/native/tile_scan.h +240 -0
  266. warp/native/tuple.h +189 -0
  267. warp/native/vec.h +2199 -0
  268. warp/native/version.h +23 -0
  269. warp/native/volume.cpp +501 -0
  270. warp/native/volume.cu +68 -0
  271. warp/native/volume.h +970 -0
  272. warp/native/volume_builder.cu +483 -0
  273. warp/native/volume_builder.h +52 -0
  274. warp/native/volume_impl.h +70 -0
  275. warp/native/warp.cpp +1143 -0
  276. warp/native/warp.cu +4604 -0
  277. warp/native/warp.h +358 -0
  278. warp/optim/__init__.py +20 -0
  279. warp/optim/adam.py +24 -0
  280. warp/optim/linear.py +35 -0
  281. warp/optim/sgd.py +24 -0
  282. warp/paddle.py +24 -0
  283. warp/py.typed +0 -0
  284. warp/render/__init__.py +22 -0
  285. warp/render/imgui_manager.py +29 -0
  286. warp/render/render_opengl.py +24 -0
  287. warp/render/render_usd.py +24 -0
  288. warp/render/utils.py +24 -0
  289. warp/sparse.py +51 -0
  290. warp/tape.py +24 -0
  291. warp/tests/__init__.py +1 -0
  292. warp/tests/__main__.py +4 -0
  293. warp/tests/assets/curlnoise_golden.npy +0 -0
  294. warp/tests/assets/mlp_golden.npy +0 -0
  295. warp/tests/assets/pixel.npy +0 -0
  296. warp/tests/assets/pnoise_golden.npy +0 -0
  297. warp/tests/assets/spiky.usd +0 -0
  298. warp/tests/assets/test_grid.nvdb +0 -0
  299. warp/tests/assets/test_index_grid.nvdb +0 -0
  300. warp/tests/assets/test_int32_grid.nvdb +0 -0
  301. warp/tests/assets/test_vec_grid.nvdb +0 -0
  302. warp/tests/assets/torus.nvdb +0 -0
  303. warp/tests/assets/torus.usda +105 -0
  304. warp/tests/aux_test_class_kernel.py +34 -0
  305. warp/tests/aux_test_compile_consts_dummy.py +18 -0
  306. warp/tests/aux_test_conditional_unequal_types_kernels.py +29 -0
  307. warp/tests/aux_test_dependent.py +29 -0
  308. warp/tests/aux_test_grad_customs.py +29 -0
  309. warp/tests/aux_test_instancing_gc.py +26 -0
  310. warp/tests/aux_test_module_aot.py +7 -0
  311. warp/tests/aux_test_module_unload.py +23 -0
  312. warp/tests/aux_test_name_clash1.py +40 -0
  313. warp/tests/aux_test_name_clash2.py +40 -0
  314. warp/tests/aux_test_reference.py +9 -0
  315. warp/tests/aux_test_reference_reference.py +8 -0
  316. warp/tests/aux_test_square.py +16 -0
  317. warp/tests/aux_test_unresolved_func.py +22 -0
  318. warp/tests/aux_test_unresolved_symbol.py +22 -0
  319. warp/tests/cuda/__init__.py +0 -0
  320. warp/tests/cuda/test_async.py +676 -0
  321. warp/tests/cuda/test_conditional_captures.py +1147 -0
  322. warp/tests/cuda/test_ipc.py +124 -0
  323. warp/tests/cuda/test_mempool.py +233 -0
  324. warp/tests/cuda/test_multigpu.py +169 -0
  325. warp/tests/cuda/test_peer.py +139 -0
  326. warp/tests/cuda/test_pinned.py +84 -0
  327. warp/tests/cuda/test_streams.py +691 -0
  328. warp/tests/geometry/__init__.py +0 -0
  329. warp/tests/geometry/test_bvh.py +335 -0
  330. warp/tests/geometry/test_hash_grid.py +259 -0
  331. warp/tests/geometry/test_marching_cubes.py +294 -0
  332. warp/tests/geometry/test_mesh.py +318 -0
  333. warp/tests/geometry/test_mesh_query_aabb.py +392 -0
  334. warp/tests/geometry/test_mesh_query_point.py +935 -0
  335. warp/tests/geometry/test_mesh_query_ray.py +323 -0
  336. warp/tests/geometry/test_volume.py +1103 -0
  337. warp/tests/geometry/test_volume_write.py +346 -0
  338. warp/tests/interop/__init__.py +0 -0
  339. warp/tests/interop/test_dlpack.py +730 -0
  340. warp/tests/interop/test_jax.py +1673 -0
  341. warp/tests/interop/test_paddle.py +800 -0
  342. warp/tests/interop/test_torch.py +1001 -0
  343. warp/tests/run_coverage_serial.py +39 -0
  344. warp/tests/test_adam.py +162 -0
  345. warp/tests/test_arithmetic.py +1096 -0
  346. warp/tests/test_array.py +3756 -0
  347. warp/tests/test_array_reduce.py +156 -0
  348. warp/tests/test_assert.py +303 -0
  349. warp/tests/test_atomic.py +336 -0
  350. warp/tests/test_atomic_bitwise.py +209 -0
  351. warp/tests/test_atomic_cas.py +312 -0
  352. warp/tests/test_bool.py +220 -0
  353. warp/tests/test_builtins_resolution.py +732 -0
  354. warp/tests/test_closest_point_edge_edge.py +327 -0
  355. warp/tests/test_codegen.py +974 -0
  356. warp/tests/test_codegen_instancing.py +1495 -0
  357. warp/tests/test_compile_consts.py +215 -0
  358. warp/tests/test_conditional.py +298 -0
  359. warp/tests/test_context.py +35 -0
  360. warp/tests/test_copy.py +319 -0
  361. warp/tests/test_ctypes.py +618 -0
  362. warp/tests/test_dense.py +73 -0
  363. warp/tests/test_devices.py +127 -0
  364. warp/tests/test_enum.py +136 -0
  365. warp/tests/test_examples.py +424 -0
  366. warp/tests/test_fabricarray.py +998 -0
  367. warp/tests/test_fast_math.py +72 -0
  368. warp/tests/test_fem.py +2204 -0
  369. warp/tests/test_fixedarray.py +229 -0
  370. warp/tests/test_fp16.py +136 -0
  371. warp/tests/test_func.py +501 -0
  372. warp/tests/test_future_annotations.py +100 -0
  373. warp/tests/test_generics.py +656 -0
  374. warp/tests/test_grad.py +893 -0
  375. warp/tests/test_grad_customs.py +339 -0
  376. warp/tests/test_grad_debug.py +341 -0
  377. warp/tests/test_implicit_init.py +411 -0
  378. warp/tests/test_import.py +45 -0
  379. warp/tests/test_indexedarray.py +1140 -0
  380. warp/tests/test_intersect.py +103 -0
  381. warp/tests/test_iter.py +76 -0
  382. warp/tests/test_large.py +177 -0
  383. warp/tests/test_launch.py +411 -0
  384. warp/tests/test_lerp.py +151 -0
  385. warp/tests/test_linear_solvers.py +223 -0
  386. warp/tests/test_lvalue.py +427 -0
  387. warp/tests/test_map.py +526 -0
  388. warp/tests/test_mat.py +3515 -0
  389. warp/tests/test_mat_assign_copy.py +178 -0
  390. warp/tests/test_mat_constructors.py +573 -0
  391. warp/tests/test_mat_lite.py +122 -0
  392. warp/tests/test_mat_scalar_ops.py +2913 -0
  393. warp/tests/test_math.py +212 -0
  394. warp/tests/test_module_aot.py +287 -0
  395. warp/tests/test_module_hashing.py +258 -0
  396. warp/tests/test_modules_lite.py +70 -0
  397. warp/tests/test_noise.py +252 -0
  398. warp/tests/test_operators.py +299 -0
  399. warp/tests/test_options.py +129 -0
  400. warp/tests/test_overwrite.py +551 -0
  401. warp/tests/test_print.py +408 -0
  402. warp/tests/test_quat.py +2653 -0
  403. warp/tests/test_quat_assign_copy.py +145 -0
  404. warp/tests/test_rand.py +339 -0
  405. warp/tests/test_reload.py +303 -0
  406. warp/tests/test_rounding.py +157 -0
  407. warp/tests/test_runlength_encode.py +196 -0
  408. warp/tests/test_scalar_ops.py +133 -0
  409. warp/tests/test_smoothstep.py +108 -0
  410. warp/tests/test_snippet.py +318 -0
  411. warp/tests/test_sparse.py +845 -0
  412. warp/tests/test_spatial.py +2859 -0
  413. warp/tests/test_spatial_assign_copy.py +160 -0
  414. warp/tests/test_special_values.py +361 -0
  415. warp/tests/test_static.py +640 -0
  416. warp/tests/test_struct.py +901 -0
  417. warp/tests/test_tape.py +242 -0
  418. warp/tests/test_transient_module.py +93 -0
  419. warp/tests/test_triangle_closest_point.py +192 -0
  420. warp/tests/test_tuple.py +361 -0
  421. warp/tests/test_types.py +615 -0
  422. warp/tests/test_utils.py +594 -0
  423. warp/tests/test_vec.py +1408 -0
  424. warp/tests/test_vec_assign_copy.py +143 -0
  425. warp/tests/test_vec_constructors.py +325 -0
  426. warp/tests/test_vec_lite.py +80 -0
  427. warp/tests/test_vec_scalar_ops.py +2327 -0
  428. warp/tests/test_verify_fp.py +100 -0
  429. warp/tests/test_version.py +75 -0
  430. warp/tests/tile/__init__.py +0 -0
  431. warp/tests/tile/test_tile.py +1519 -0
  432. warp/tests/tile/test_tile_atomic_bitwise.py +403 -0
  433. warp/tests/tile/test_tile_cholesky.py +608 -0
  434. warp/tests/tile/test_tile_load.py +724 -0
  435. warp/tests/tile/test_tile_mathdx.py +156 -0
  436. warp/tests/tile/test_tile_matmul.py +179 -0
  437. warp/tests/tile/test_tile_mlp.py +400 -0
  438. warp/tests/tile/test_tile_reduce.py +950 -0
  439. warp/tests/tile/test_tile_shared_memory.py +376 -0
  440. warp/tests/tile/test_tile_sort.py +121 -0
  441. warp/tests/tile/test_tile_view.py +173 -0
  442. warp/tests/unittest_serial.py +47 -0
  443. warp/tests/unittest_suites.py +430 -0
  444. warp/tests/unittest_utils.py +469 -0
  445. warp/tests/walkthrough_debug.py +95 -0
  446. warp/torch.py +24 -0
  447. warp/types.py +51 -0
  448. warp/utils.py +31 -0
  449. warp_lang-1.10.0.dist-info/METADATA +459 -0
  450. warp_lang-1.10.0.dist-info/RECORD +468 -0
  451. warp_lang-1.10.0.dist-info/WHEEL +5 -0
  452. warp_lang-1.10.0.dist-info/licenses/LICENSE.md +176 -0
  453. warp_lang-1.10.0.dist-info/licenses/licenses/Gaia-LICENSE.txt +6 -0
  454. warp_lang-1.10.0.dist-info/licenses/licenses/appdirs-LICENSE.txt +22 -0
  455. warp_lang-1.10.0.dist-info/licenses/licenses/asset_pixel_jpg-LICENSE.txt +3 -0
  456. warp_lang-1.10.0.dist-info/licenses/licenses/cuda-LICENSE.txt +1582 -0
  457. warp_lang-1.10.0.dist-info/licenses/licenses/dlpack-LICENSE.txt +201 -0
  458. warp_lang-1.10.0.dist-info/licenses/licenses/fp16-LICENSE.txt +28 -0
  459. warp_lang-1.10.0.dist-info/licenses/licenses/libmathdx-LICENSE.txt +220 -0
  460. warp_lang-1.10.0.dist-info/licenses/licenses/llvm-LICENSE.txt +279 -0
  461. warp_lang-1.10.0.dist-info/licenses/licenses/moller-LICENSE.txt +16 -0
  462. warp_lang-1.10.0.dist-info/licenses/licenses/nanovdb-LICENSE.txt +2 -0
  463. warp_lang-1.10.0.dist-info/licenses/licenses/nvrtc-LICENSE.txt +1592 -0
  464. warp_lang-1.10.0.dist-info/licenses/licenses/svd-LICENSE.txt +23 -0
  465. warp_lang-1.10.0.dist-info/licenses/licenses/unittest_parallel-LICENSE.txt +21 -0
  466. warp_lang-1.10.0.dist-info/licenses/licenses/usd-LICENSE.txt +213 -0
  467. warp_lang-1.10.0.dist-info/licenses/licenses/windingnumber-LICENSE.txt +21 -0
  468. warp_lang-1.10.0.dist-info/top_level.txt +1 -0
@@ -0,0 +1,800 @@
1
+ # SPDX-FileCopyrightText: Copyright (c) 2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2
+ # SPDX-License-Identifier: Apache-2.0
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+
16
+ import unittest
17
+
18
+ import warp as wp
19
+ from warp.tests.unittest_utils import *
20
+
21
+
22
+ @wp.kernel
23
+ def op_kernel(x: wp.array(dtype=float), y: wp.array(dtype=float)):
24
+ tid = wp.tid()
25
+ y[tid] = 0.5 - x[tid] * 2.0
26
+
27
+
28
+ @wp.kernel
29
+ def inc(a: wp.array(dtype=float)):
30
+ tid = wp.tid()
31
+ a[tid] = a[tid] + 1.0
32
+
33
+
34
+ @wp.kernel
35
+ def inc_vector(a: wp.array(dtype=wp.vec3f)):
36
+ tid = wp.tid()
37
+ a[tid] = a[tid] + wp.vec3f(1.0)
38
+
39
+
40
+ @wp.kernel
41
+ def inc_matrix(a: wp.array(dtype=wp.mat22f)):
42
+ tid = wp.tid()
43
+ a[tid] = a[tid] + wp.mat22f(1.0)
44
+
45
+
46
+ @wp.kernel
47
+ def arange(start: int, step: int, a: wp.array(dtype=int)):
48
+ tid = wp.tid()
49
+ a[tid] = start + step * tid
50
+
51
+
52
+ # copy elements between non-contiguous 1d arrays of float
53
+ @wp.kernel
54
+ def copy1d_float_kernel(dst: wp.array(dtype=float), src: wp.array(dtype=float)):
55
+ i = wp.tid()
56
+ dst[i] = src[i]
57
+
58
+
59
+ # copy elements between non-contiguous 2d arrays of float
60
+ @wp.kernel
61
+ def copy2d_float_kernel(dst: wp.array2d(dtype=float), src: wp.array2d(dtype=float)):
62
+ i, j = wp.tid()
63
+ dst[i, j] = src[i, j]
64
+
65
+
66
+ # copy elements between non-contiguous 3d arrays of float
67
+ @wp.kernel
68
+ def copy3d_float_kernel(dst: wp.array3d(dtype=float), src: wp.array3d(dtype=float)):
69
+ i, j, k = wp.tid()
70
+ dst[i, j, k] = src[i, j, k]
71
+
72
+
73
+ # copy elements between non-contiguous 2d arrays of vec3
74
+ @wp.kernel
75
+ def copy2d_vec3_kernel(dst: wp.array2d(dtype=wp.vec3), src: wp.array2d(dtype=wp.vec3)):
76
+ i, j = wp.tid()
77
+ dst[i, j] = src[i, j]
78
+
79
+
80
+ # copy elements between non-contiguous 2d arrays of mat22
81
+ @wp.kernel
82
+ def copy2d_mat22_kernel(dst: wp.array2d(dtype=wp.mat22), src: wp.array2d(dtype=wp.mat22)):
83
+ i, j = wp.tid()
84
+ dst[i, j] = src[i, j]
85
+
86
+
87
+ def test_dtype_from_paddle(test, device):
88
+ import paddle
89
+
90
+ def test_conversions(paddle_type, warp_type):
91
+ test.assertEqual(wp.dtype_from_paddle(paddle_type), warp_type)
92
+
93
+ test_conversions(paddle.float16, wp.float16)
94
+ test_conversions(paddle.float32, wp.float32)
95
+ test_conversions(paddle.float64, wp.float64)
96
+ test_conversions(paddle.int8, wp.int8)
97
+ test_conversions(paddle.int16, wp.int16)
98
+ test_conversions(paddle.int32, wp.int32)
99
+ test_conversions(paddle.int64, wp.int64)
100
+ test_conversions(paddle.uint8, wp.uint8)
101
+ test_conversions(paddle.bool, wp.bool)
102
+
103
+
104
+ def test_dtype_to_paddle(test, device):
105
+ import paddle
106
+
107
+ def test_conversions(warp_type, paddle_type):
108
+ test.assertEqual(wp.dtype_to_paddle(warp_type), paddle_type)
109
+
110
+ test_conversions(wp.float16, paddle.float16)
111
+ test_conversions(wp.float32, paddle.float32)
112
+ test_conversions(wp.float64, paddle.float64)
113
+ test_conversions(wp.int8, paddle.int8)
114
+ test_conversions(wp.int16, paddle.int16)
115
+ test_conversions(wp.int32, paddle.int32)
116
+ test_conversions(wp.int64, paddle.int64)
117
+ test_conversions(wp.uint8, paddle.uint8)
118
+ test_conversions(wp.uint16, paddle.int16)
119
+ test_conversions(wp.uint32, paddle.int32)
120
+ test_conversions(wp.uint64, paddle.int64)
121
+ test_conversions(wp.bool, paddle.bool)
122
+
123
+
124
+ def test_device_conversion(test, device):
125
+ paddle_device = wp.device_to_paddle(device)
126
+ warp_device = wp.device_from_paddle(paddle_device)
127
+ test.assertEqual(warp_device, device)
128
+
129
+
130
+ def test_paddle_zerocopy(test, device):
131
+ import paddle
132
+
133
+ a = wp.zeros(10, dtype=wp.float32, device=device)
134
+ t = wp.to_paddle(a)
135
+ assert a.ptr == t.data_ptr()
136
+
137
+ paddle_device = wp.device_to_paddle(device)
138
+
139
+ t = paddle.zeros([10], dtype=paddle.float32).to(device=paddle_device)
140
+ a = wp.from_paddle(t)
141
+ assert a.ptr == t.data_ptr()
142
+
143
+
144
+ def test_from_paddle(test, device):
145
+ import paddle
146
+
147
+ paddle_device = wp.device_to_paddle(device)
148
+
149
+ # automatically determine warp dtype
150
+ def wrap_scalar_tensor_implicit(paddle_dtype, expected_warp_dtype):
151
+ t = paddle.zeros([10], dtype=paddle_dtype).to(device=paddle_device)
152
+ a = wp.from_paddle(t)
153
+ assert a.dtype == expected_warp_dtype
154
+ assert a.shape == tuple(t.shape)
155
+
156
+ wrap_scalar_tensor_implicit(paddle.float64, wp.float64)
157
+ wrap_scalar_tensor_implicit(paddle.float32, wp.float32)
158
+ wrap_scalar_tensor_implicit(paddle.float16, wp.float16)
159
+ wrap_scalar_tensor_implicit(paddle.int64, wp.int64)
160
+ wrap_scalar_tensor_implicit(paddle.int32, wp.int32)
161
+ wrap_scalar_tensor_implicit(paddle.int16, wp.int16)
162
+ wrap_scalar_tensor_implicit(paddle.int8, wp.int8)
163
+ wrap_scalar_tensor_implicit(paddle.uint8, wp.uint8)
164
+ wrap_scalar_tensor_implicit(paddle.bool, wp.bool)
165
+
166
+ # explicitly specify warp dtype
167
+ def wrap_scalar_tensor_explicit(paddle_dtype, expected_warp_dtype):
168
+ t = paddle.zeros([10], dtype=paddle_dtype).to(device=paddle_device)
169
+ a = wp.from_paddle(t, expected_warp_dtype)
170
+ assert a.dtype == expected_warp_dtype
171
+ assert a.shape == tuple(t.shape)
172
+
173
+ wrap_scalar_tensor_explicit(paddle.float64, wp.float64)
174
+ wrap_scalar_tensor_explicit(paddle.float32, wp.float32)
175
+ wrap_scalar_tensor_explicit(paddle.float16, wp.float16)
176
+ wrap_scalar_tensor_explicit(paddle.int64, wp.int64)
177
+ wrap_scalar_tensor_explicit(paddle.int64, wp.uint64)
178
+ wrap_scalar_tensor_explicit(paddle.int32, wp.int32)
179
+ wrap_scalar_tensor_explicit(paddle.int32, wp.uint32)
180
+ wrap_scalar_tensor_explicit(paddle.int16, wp.int16)
181
+ wrap_scalar_tensor_explicit(paddle.int16, wp.uint16)
182
+ wrap_scalar_tensor_explicit(paddle.int8, wp.int8)
183
+ wrap_scalar_tensor_explicit(paddle.int8, wp.uint8)
184
+ wrap_scalar_tensor_explicit(paddle.uint8, wp.uint8)
185
+ wrap_scalar_tensor_explicit(paddle.uint8, wp.int8)
186
+ wrap_scalar_tensor_explicit(paddle.bool, wp.uint8)
187
+ wrap_scalar_tensor_explicit(paddle.bool, wp.int8)
188
+ wrap_scalar_tensor_explicit(paddle.bool, wp.bool)
189
+
190
+ def wrap_vec_tensor(n, desired_warp_dtype):
191
+ t = paddle.zeros((10, n), dtype=paddle.float32).to(device=paddle_device)
192
+ a = wp.from_paddle(t, desired_warp_dtype)
193
+ assert a.dtype == desired_warp_dtype
194
+ assert a.shape == (10,)
195
+
196
+ wrap_vec_tensor(2, wp.vec2)
197
+ wrap_vec_tensor(3, wp.vec3)
198
+ wrap_vec_tensor(4, wp.vec4)
199
+ wrap_vec_tensor(6, wp.spatial_vector)
200
+ wrap_vec_tensor(7, wp.transform)
201
+
202
+ def wrap_mat_tensor(n, m, desired_warp_dtype):
203
+ t = paddle.zeros((10, n, m), dtype=paddle.float32).to(device=paddle_device)
204
+ a = wp.from_paddle(t, desired_warp_dtype)
205
+ assert a.dtype == desired_warp_dtype
206
+ assert a.shape == (10,)
207
+
208
+ wrap_mat_tensor(2, 2, wp.mat22)
209
+ wrap_mat_tensor(3, 3, wp.mat33)
210
+ wrap_mat_tensor(4, 4, wp.mat44)
211
+ wrap_mat_tensor(6, 6, wp.spatial_matrix)
212
+
213
+ def wrap_vec_tensor_with_grad(n, desired_warp_dtype):
214
+ t = paddle.zeros((10, n), dtype=paddle.float32).to(device=paddle_device)
215
+ a = wp.from_paddle(t, desired_warp_dtype)
216
+ a.requires_grad = True
217
+ assert a.dtype == desired_warp_dtype
218
+ assert a.shape == (10,)
219
+
220
+ wrap_vec_tensor_with_grad(2, wp.vec2)
221
+ wrap_vec_tensor_with_grad(3, wp.vec3)
222
+ wrap_vec_tensor_with_grad(4, wp.vec4)
223
+ wrap_vec_tensor_with_grad(6, wp.spatial_vector)
224
+ wrap_vec_tensor_with_grad(7, wp.transform)
225
+
226
+ def wrap_mat_tensor_with_grad(n, m, desired_warp_dtype):
227
+ t = paddle.zeros((10, n, m), dtype=paddle.float32).to(device=paddle_device)
228
+ a = wp.from_paddle(t, desired_warp_dtype, requires_grad=True)
229
+ assert a.dtype == desired_warp_dtype
230
+ assert a.shape == (10,)
231
+
232
+ wrap_mat_tensor_with_grad(2, 2, wp.mat22)
233
+ wrap_mat_tensor_with_grad(3, 3, wp.mat33)
234
+ wrap_mat_tensor_with_grad(4, 4, wp.mat44)
235
+ wrap_mat_tensor_with_grad(6, 6, wp.spatial_matrix)
236
+
237
+
238
+ def test_array_ctype_from_paddle(test, device):
239
+ import paddle
240
+
241
+ paddle_device = wp.device_to_paddle(device)
242
+
243
+ # automatically determine warp dtype
244
+ def wrap_scalar_tensor_implicit(paddle_dtype):
245
+ t = paddle.zeros([10], dtype=paddle_dtype).to(device=paddle_device)
246
+ a = wp.from_paddle(t, return_ctype=True)
247
+ warp_dtype = wp.dtype_from_paddle(paddle_dtype)
248
+ ctype_size = ctypes.sizeof(warp_dtype._type_)
249
+ assert a.data == t.data_ptr()
250
+ assert a.grad == 0
251
+ assert a.ndim == 1
252
+ assert a.shape[0] == t.shape[0]
253
+ assert a.strides[0] == t.strides[0] * ctype_size
254
+
255
+ wrap_scalar_tensor_implicit(paddle.float64)
256
+ wrap_scalar_tensor_implicit(paddle.float32)
257
+ wrap_scalar_tensor_implicit(paddle.float16)
258
+ wrap_scalar_tensor_implicit(paddle.int64)
259
+ wrap_scalar_tensor_implicit(paddle.int32)
260
+ wrap_scalar_tensor_implicit(paddle.int16)
261
+ wrap_scalar_tensor_implicit(paddle.int8)
262
+ wrap_scalar_tensor_implicit(paddle.uint8)
263
+ wrap_scalar_tensor_implicit(paddle.bool)
264
+
265
+ # explicitly specify warp dtype
266
+ def wrap_scalar_tensor_explicit(paddle_dtype, warp_dtype):
267
+ t = paddle.zeros([10], dtype=paddle_dtype).to(device=paddle_device)
268
+ a = wp.from_paddle(t, dtype=warp_dtype, return_ctype=True)
269
+ ctype_size = ctypes.sizeof(warp_dtype._type_)
270
+ assert a.data == t.data_ptr()
271
+ assert a.grad == 0
272
+ assert a.ndim == 1
273
+ assert a.shape[0] == t.shape[0]
274
+ assert a.strides[0] == t.strides[0] * ctype_size
275
+
276
+ wrap_scalar_tensor_explicit(paddle.float64, wp.float64)
277
+ wrap_scalar_tensor_explicit(paddle.float32, wp.float32)
278
+ wrap_scalar_tensor_explicit(paddle.float16, wp.float16)
279
+ wrap_scalar_tensor_explicit(paddle.int64, wp.int64)
280
+ wrap_scalar_tensor_explicit(paddle.int64, wp.uint64)
281
+ wrap_scalar_tensor_explicit(paddle.int32, wp.int32)
282
+ wrap_scalar_tensor_explicit(paddle.int32, wp.uint32)
283
+ wrap_scalar_tensor_explicit(paddle.int16, wp.int16)
284
+ wrap_scalar_tensor_explicit(paddle.int16, wp.uint16)
285
+ wrap_scalar_tensor_explicit(paddle.int8, wp.int8)
286
+ wrap_scalar_tensor_explicit(paddle.int8, wp.uint8)
287
+ wrap_scalar_tensor_explicit(paddle.uint8, wp.uint8)
288
+ wrap_scalar_tensor_explicit(paddle.uint8, wp.int8)
289
+ wrap_scalar_tensor_explicit(paddle.bool, wp.uint8)
290
+ wrap_scalar_tensor_explicit(paddle.bool, wp.int8)
291
+ wrap_scalar_tensor_explicit(paddle.bool, wp.bool)
292
+
293
+ def wrap_vec_tensor(vec_dtype):
294
+ t = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
295
+ a = wp.from_paddle(t, dtype=vec_dtype, return_ctype=True)
296
+ ctype_size = ctypes.sizeof(vec_dtype._type_)
297
+ assert a.data == t.data_ptr()
298
+ assert a.grad == 0
299
+ assert a.ndim == 1
300
+ assert a.shape[0] == t.shape[0]
301
+ assert a.strides[0] == t.strides[0] * ctype_size
302
+
303
+ wrap_vec_tensor(wp.vec2)
304
+ wrap_vec_tensor(wp.vec3)
305
+ wrap_vec_tensor(wp.vec4)
306
+ wrap_vec_tensor(wp.spatial_vector)
307
+ wrap_vec_tensor(wp.transform)
308
+
309
+ def wrap_mat_tensor(mat_dtype):
310
+ t = paddle.zeros((10, *mat_dtype._shape_), dtype=paddle.float32).to(device=paddle_device)
311
+ a = wp.from_paddle(t, dtype=mat_dtype, return_ctype=True)
312
+ ctype_size = ctypes.sizeof(mat_dtype._type_)
313
+ assert a.data == t.data_ptr()
314
+ assert a.grad == 0
315
+ assert a.ndim == 1
316
+ assert a.shape[0] == t.shape[0]
317
+ assert a.strides[0] == t.strides[0] * ctype_size
318
+
319
+ wrap_mat_tensor(wp.mat22)
320
+ wrap_mat_tensor(wp.mat33)
321
+ wrap_mat_tensor(wp.mat44)
322
+ wrap_mat_tensor(wp.spatial_matrix)
323
+
324
+ def wrap_vec_tensor_with_existing_grad(vec_dtype):
325
+ t = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
326
+ t.stop_gradient = False
327
+ t.grad_ = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
328
+ a = wp.from_paddle(t, dtype=vec_dtype, return_ctype=True)
329
+ ctype_size = ctypes.sizeof(vec_dtype._type_)
330
+ assert a.data == t.data_ptr()
331
+ assert a.grad == t.grad.data_ptr()
332
+ assert a.ndim == 1
333
+ assert a.shape[0] == t.shape[0]
334
+ assert a.strides[0] == t.strides[0] * ctype_size
335
+
336
+ wrap_vec_tensor_with_existing_grad(wp.vec2)
337
+ wrap_vec_tensor_with_existing_grad(wp.vec3)
338
+ wrap_vec_tensor_with_existing_grad(wp.vec4)
339
+ wrap_vec_tensor_with_existing_grad(wp.spatial_vector)
340
+ wrap_vec_tensor_with_existing_grad(wp.transform)
341
+
342
+ def wrap_vec_tensor_with_new_grad(vec_dtype):
343
+ t = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
344
+ a = wp.from_paddle(t, dtype=vec_dtype, requires_grad=True, return_ctype=True)
345
+ ctype_size = ctypes.sizeof(vec_dtype._type_)
346
+ assert a.data == t.data_ptr()
347
+ assert a.grad == t.grad.data_ptr()
348
+ assert a.ndim == 1
349
+ assert a.shape[0] == t.shape[0]
350
+ assert a.strides[0] == t.strides[0] * ctype_size
351
+
352
+ wrap_vec_tensor_with_new_grad(wp.vec2)
353
+ wrap_vec_tensor_with_new_grad(wp.vec3)
354
+ wrap_vec_tensor_with_new_grad(wp.vec4)
355
+ wrap_vec_tensor_with_new_grad(wp.spatial_vector)
356
+ wrap_vec_tensor_with_new_grad(wp.transform)
357
+
358
+ def wrap_vec_tensor_with_paddle_grad(vec_dtype):
359
+ t = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
360
+ grad = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
361
+ a = wp.from_paddle(t, dtype=vec_dtype, grad=grad, return_ctype=True)
362
+ ctype_size = ctypes.sizeof(vec_dtype._type_)
363
+ assert a.data == t.data_ptr()
364
+ assert a.grad == grad.data_ptr()
365
+ assert a.ndim == 1
366
+ assert a.shape[0] == t.shape[0]
367
+ assert a.strides[0] == t.strides[0] * ctype_size
368
+
369
+ wrap_vec_tensor_with_paddle_grad(wp.vec2)
370
+ wrap_vec_tensor_with_paddle_grad(wp.vec3)
371
+ wrap_vec_tensor_with_paddle_grad(wp.vec4)
372
+ wrap_vec_tensor_with_paddle_grad(wp.spatial_vector)
373
+ wrap_vec_tensor_with_paddle_grad(wp.transform)
374
+
375
+ def wrap_vec_tensor_with_warp_grad(vec_dtype):
376
+ t = paddle.zeros((10, vec_dtype._length_), dtype=paddle.float32).to(device=paddle_device)
377
+ grad = wp.zeros(10, dtype=vec_dtype, device=device)
378
+ a = wp.from_paddle(t, dtype=vec_dtype, grad=grad, return_ctype=True)
379
+ ctype_size = ctypes.sizeof(vec_dtype._type_)
380
+ assert a.data == t.data_ptr()
381
+ assert a.grad == grad.ptr
382
+ assert a.ndim == 1
383
+ assert a.shape[0] == t.shape[0]
384
+ assert a.strides[0] == t.strides[0] * ctype_size
385
+
386
+ wrap_vec_tensor_with_warp_grad(wp.vec2)
387
+ wrap_vec_tensor_with_warp_grad(wp.vec3)
388
+ wrap_vec_tensor_with_warp_grad(wp.vec4)
389
+ wrap_vec_tensor_with_warp_grad(wp.spatial_vector)
390
+ wrap_vec_tensor_with_warp_grad(wp.transform)
391
+
392
+
393
+ def test_to_paddle(test, device):
394
+ import paddle
395
+
396
+ def wrap_scalar_array(warp_dtype, expected_paddle_dtype):
397
+ a = wp.zeros(10, dtype=warp_dtype, device=device)
398
+ t = wp.to_paddle(a)
399
+ assert t.dtype == expected_paddle_dtype
400
+ assert tuple(t.shape) == a.shape
401
+
402
+ wrap_scalar_array(wp.float64, paddle.float64)
403
+ wrap_scalar_array(wp.float32, paddle.float32)
404
+ wrap_scalar_array(wp.float16, paddle.float16)
405
+ wrap_scalar_array(wp.int64, paddle.int64)
406
+ wrap_scalar_array(wp.int32, paddle.int32)
407
+ wrap_scalar_array(wp.int16, paddle.int16)
408
+ wrap_scalar_array(wp.int8, paddle.int8)
409
+ wrap_scalar_array(wp.uint8, paddle.uint8)
410
+ wrap_scalar_array(wp.bool, paddle.bool)
411
+
412
+ # not supported by paddle
413
+ # wrap_scalar_array(wp.uint64, paddle.int64)
414
+ # wrap_scalar_array(wp.uint32, paddle.int32)
415
+ # wrap_scalar_array(wp.uint16, paddle.int16)
416
+
417
+ def wrap_vec_array(n, warp_dtype):
418
+ a = wp.zeros(10, dtype=warp_dtype, device=device)
419
+ t = wp.to_paddle(a)
420
+ assert t.dtype == paddle.float32
421
+ assert tuple(t.shape) == (10, n)
422
+
423
+ wrap_vec_array(2, wp.vec2)
424
+ wrap_vec_array(3, wp.vec3)
425
+ wrap_vec_array(4, wp.vec4)
426
+ wrap_vec_array(6, wp.spatial_vector)
427
+ wrap_vec_array(7, wp.transform)
428
+
429
+ def wrap_mat_array(n, m, warp_dtype):
430
+ a = wp.zeros(10, dtype=warp_dtype, device=device)
431
+ t = wp.to_paddle(a)
432
+ assert t.dtype == paddle.float32
433
+ assert tuple(t.shape) == (10, n, m)
434
+
435
+ wrap_mat_array(2, 2, wp.mat22)
436
+ wrap_mat_array(3, 3, wp.mat33)
437
+ wrap_mat_array(4, 4, wp.mat44)
438
+ wrap_mat_array(6, 6, wp.spatial_matrix)
439
+
440
+
441
+ def test_from_paddle_slices(test, device):
442
+ import paddle
443
+
444
+ paddle_device = wp.device_to_paddle(device)
445
+
446
+ # 1D slice, contiguous
447
+ t_base = paddle.arange(10, dtype=paddle.float32).to(device=paddle_device)
448
+ t = t_base[2:9]
449
+ a = wp.from_paddle(t)
450
+ assert a.ptr == t.data_ptr()
451
+ assert a.is_contiguous
452
+ assert a.shape == tuple(t.shape)
453
+ assert_np_equal(a.numpy(), t.numpy())
454
+
455
+ # 1D slice with non-contiguous stride
456
+ t_base = paddle.arange(10, dtype=paddle.float32).to(device=paddle_device)
457
+ t = t_base[2:9:2]
458
+ a = wp.from_paddle(t)
459
+ assert a.ptr == t.data_ptr()
460
+ assert not a.is_contiguous
461
+ assert a.shape == tuple(t.shape)
462
+ # copy contents to contiguous array
463
+ a_contiguous = wp.empty_like(a)
464
+ wp.launch(copy1d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
465
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
466
+
467
+ # 2D slices (non-contiguous)
468
+ t_base = paddle.arange(24, dtype=paddle.float32).to(device=paddle_device).reshape((4, 6))
469
+ t = t_base[1:3, 2:5]
470
+ a = wp.from_paddle(t)
471
+ assert a.ptr == t.data_ptr()
472
+ assert not a.is_contiguous
473
+ assert a.shape == tuple(t.shape)
474
+ # copy contents to contiguous array
475
+ a_contiguous = wp.empty_like(a)
476
+ wp.launch(copy2d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
477
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
478
+
479
+ # 3D slices (non-contiguous)
480
+ t_base = paddle.arange(36, dtype=paddle.float32).to(device=paddle_device).reshape((4, 3, 3))
481
+ t = t_base[::2, 0:1, 1:2]
482
+ a = wp.from_paddle(t)
483
+ assert a.ptr == t.data_ptr()
484
+ assert not a.is_contiguous
485
+ assert a.shape == tuple(t.shape)
486
+ # copy contents to contiguous array
487
+ a_contiguous = wp.empty_like(a)
488
+ wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
489
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
490
+
491
+ # 2D slices of vec3 (inner contiguous, outer non-contiguous)
492
+ t_base = paddle.arange(150, dtype=paddle.float32).to(device=paddle_device).reshape((10, 5, 3))
493
+ t = t_base[1:7:2, 2:5]
494
+ a = wp.from_paddle(t, dtype=wp.vec3)
495
+ assert a.ptr == t.data_ptr()
496
+ assert not a.is_contiguous
497
+ assert a.shape == tuple(t.shape[:-1])
498
+ # copy contents to contiguous array
499
+ a_contiguous = wp.empty_like(a)
500
+ wp.launch(copy2d_vec3_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
501
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
502
+
503
+ # 2D slices of mat22 (inner contiguous, outer non-contiguous)
504
+ t_base = paddle.arange(200, dtype=paddle.float32).to(device=paddle_device).reshape((10, 5, 2, 2))
505
+ t = t_base[1:7:2, 2:5]
506
+ a = wp.from_paddle(t, dtype=wp.mat22)
507
+ assert a.ptr == t.data_ptr()
508
+ assert not a.is_contiguous
509
+ assert a.shape == tuple(t.shape[:-2])
510
+ # copy contents to contiguous array
511
+ a_contiguous = wp.empty_like(a)
512
+ wp.launch(copy2d_mat22_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
513
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
514
+
515
+
516
+ def test_from_paddle_zero_strides(test, device):
517
+ import paddle
518
+
519
+ paddle_device = wp.device_to_paddle(device)
520
+
521
+ t_base = paddle.arange(9, dtype=paddle.float32).to(device=paddle_device).reshape((3, 3))
522
+
523
+ # expand outermost dimension
524
+ t = t_base.unsqueeze(0).expand([3, -1, -1])
525
+ a = wp.from_paddle(t)
526
+ assert a.ptr == t.data_ptr()
527
+ assert a.is_contiguous
528
+ assert a.shape == tuple(t.shape)
529
+ a_contiguous = wp.empty_like(a)
530
+ wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
531
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
532
+
533
+ # expand middle dimension
534
+ t = t_base.unsqueeze(1).expand([-1, 3, -1])
535
+ a = wp.from_paddle(t)
536
+ assert a.ptr == t.data_ptr()
537
+ assert a.is_contiguous
538
+ assert a.shape == tuple(t.shape)
539
+ a_contiguous = wp.empty_like(a)
540
+ wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
541
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
542
+
543
+ # expand innermost dimension
544
+ t = t_base.unsqueeze(2).expand([-1, -1, 3])
545
+ a = wp.from_paddle(t)
546
+ assert a.ptr == t.data_ptr()
547
+ assert a.is_contiguous
548
+ assert a.shape == tuple(t.shape)
549
+ a_contiguous = wp.empty_like(a)
550
+ wp.launch(copy3d_float_kernel, dim=a.shape, inputs=[a_contiguous, a], device=device)
551
+ assert_np_equal(a_contiguous.numpy(), t.numpy())
552
+
553
+
554
+ def test_paddle_autograd(test, device):
555
+ """Test paddle autograd with a custom Warp op"""
556
+
557
+ import paddle
558
+
559
+ # custom autograd op
560
+ class TestFunc(paddle.autograd.PyLayer):
561
+ @staticmethod
562
+ def forward(ctx, x):
563
+ # ensure Paddle operations complete before running Warp
564
+ wp.synchronize_device()
565
+
566
+ # allocate output array
567
+ y = paddle.empty_like(x)
568
+
569
+ ctx.x = x
570
+ ctx.y = y
571
+
572
+ wp.launch(kernel=op_kernel, dim=len(x), inputs=[wp.from_paddle(x)], outputs=[wp.from_paddle(y)])
573
+
574
+ # ensure Warp operations complete before returning data to Paddle
575
+ wp.synchronize_device()
576
+
577
+ return y
578
+
579
+ @staticmethod
580
+ def backward(ctx, adj_y):
581
+ # ensure Paddle operations complete before running Warp
582
+ wp.synchronize_device()
583
+
584
+ # adjoints should be allocated as zero initialized
585
+ adj_x = paddle.zeros_like(ctx.x).contiguous()
586
+ adj_y = adj_y.contiguous()
587
+
588
+ wp_x = wp.from_paddle(ctx.x, grad=adj_x)
589
+ wp_y = wp.from_paddle(ctx.y, grad=adj_y)
590
+
591
+ wp.launch(
592
+ kernel=op_kernel,
593
+ dim=len(ctx.x),
594
+ # fwd inputs
595
+ inputs=[wp_x],
596
+ outputs=[wp_y],
597
+ # adj inputs (already stored in input/output arrays, passing null pointers)
598
+ adj_inputs=[None],
599
+ adj_outputs=[None],
600
+ adjoint=True,
601
+ )
602
+
603
+ # ensure Warp operations complete before returning data to Paddle
604
+ wp.synchronize_device()
605
+
606
+ return adj_x
607
+
608
+ # run autograd on given device
609
+ with wp.ScopedDevice(device):
610
+ paddle_device = wp.device_to_paddle(device)
611
+
612
+ # input data
613
+ x = paddle.ones(16, dtype=paddle.float32).to(device=paddle_device)
614
+ x.stop_gradient = False
615
+
616
+ # execute op
617
+ y = TestFunc.apply(x)
618
+
619
+ # compute grads
620
+ l = y.sum()
621
+ l.backward()
622
+
623
+ passed = (x.grad == -2.0).all()
624
+ assert passed.item()
625
+
626
+
627
+ def test_warp_graph_warp_stream(test, device):
628
+ """Capture Warp graph on Warp stream"""
629
+
630
+ import paddle
631
+
632
+ paddle_device = wp.device_to_paddle(device)
633
+
634
+ n = 1024 * 1024
635
+ t = paddle.zeros(n, dtype=paddle.float32).to(device=paddle_device)
636
+ a = wp.from_paddle(t)
637
+
638
+ # make paddle use the warp stream from the given device
639
+ paddle_stream = wp.stream_to_paddle(device)
640
+
641
+ # capture graph
642
+ with wp.ScopedDevice(device), paddle.device.stream_guard(paddle.device.Stream(paddle_stream)):
643
+ wp.capture_begin(force_module_load=False)
644
+ try:
645
+ t += 1.0
646
+ wp.launch(inc, dim=n, inputs=[a])
647
+ t += 1.0
648
+ wp.launch(inc, dim=n, inputs=[a])
649
+ finally:
650
+ g = wp.capture_end()
651
+
652
+ # replay graph
653
+ num_iters = 10
654
+ for _i in range(num_iters):
655
+ wp.capture_launch(g)
656
+
657
+ passed = (t == num_iters * 4.0).all()
658
+ assert passed.item()
659
+
660
+
661
+ def test_warp_graph_paddle_stream(test, device):
662
+ """Capture Warp graph on Paddle stream"""
663
+
664
+ wp.load_module(device=device)
665
+
666
+ import paddle
667
+
668
+ paddle_device = wp.device_to_paddle(device)
669
+
670
+ n = 1024 * 1024
671
+ t = paddle.zeros(n, dtype=paddle.float32).to(device=paddle_device)
672
+ a = wp.from_paddle(t)
673
+
674
+ # create a device-specific paddle stream to use for capture
675
+ # (the default paddle stream is not suitable for graph capture)
676
+ paddle_stream = paddle.device.Stream(device=paddle_device)
677
+
678
+ # make warp use the same stream
679
+ warp_stream = wp.stream_from_paddle(paddle_stream)
680
+
681
+ # capture graph
682
+ with wp.ScopedStream(warp_stream):
683
+ wp.capture_begin(force_module_load=False)
684
+ try:
685
+ t += 1.0
686
+ wp.launch(inc, dim=n, inputs=[a])
687
+ t += 1.0
688
+ wp.launch(inc, dim=n, inputs=[a])
689
+ finally:
690
+ g = wp.capture_end()
691
+
692
+ # replay graph
693
+ num_iters = 10
694
+ for _i in range(num_iters):
695
+ wp.capture_launch(g)
696
+
697
+ passed = (t == num_iters * 4.0).all()
698
+ assert passed.item()
699
+
700
+
701
+ def test_direct(test, device):
702
+ """Pass Paddle tensors to Warp kernels directly"""
703
+
704
+ import paddle
705
+
706
+ paddle_device = wp.device_to_paddle(device)
707
+ n = 12
708
+
709
+ s = paddle.arange(n, dtype=paddle.float32).to(device=paddle_device)
710
+ v = paddle.arange(n, dtype=paddle.float32).to(device=paddle_device).reshape((n // 3, 3))
711
+ m = paddle.arange(n, dtype=paddle.float32).to(device=paddle_device).reshape((n // 4, 2, 2))
712
+
713
+ wp.launch(inc, dim=n, inputs=[s], device=device)
714
+ wp.launch(inc_vector, dim=n // 3, inputs=[v], device=device)
715
+ wp.launch(inc_matrix, dim=n // 4, inputs=[m], device=device)
716
+
717
+ expected = paddle.arange(1, n + 1, dtype=paddle.float32).to(device=paddle_device)
718
+
719
+ assert paddle.equal_all(s, expected).item()
720
+ assert paddle.equal_all(v.reshape([n]), expected).item()
721
+ assert paddle.equal_all(m.reshape([n]), expected).item()
722
+
723
+
724
+ class TestPaddle(unittest.TestCase):
725
+ pass
726
+
727
+
728
+ test_devices = get_test_devices()
729
+
730
+ try:
731
+ import paddle
732
+
733
+ # check which Warp devices work with Paddle
734
+ # CUDA devices may fail if Paddle was not compiled with CUDA support
735
+ paddle_compatible_devices = []
736
+ paddle_compatible_cuda_devices = []
737
+
738
+ for d in test_devices:
739
+ try:
740
+ t = paddle.arange(10).to(device=wp.device_to_paddle(d))
741
+ t += 1
742
+ paddle_compatible_devices.append(d)
743
+ if d.is_cuda:
744
+ paddle_compatible_cuda_devices.append(d)
745
+ except Exception as e:
746
+ print(f"Skipping Paddle tests on device '{d}' due to exception: {e}")
747
+
748
+ add_function_test(TestPaddle, "test_dtype_from_paddle", test_dtype_from_paddle, devices=None)
749
+ add_function_test(TestPaddle, "test_dtype_to_paddle", test_dtype_to_paddle, devices=None)
750
+
751
+ if paddle_compatible_devices:
752
+ add_function_test(
753
+ TestPaddle, "test_device_conversion", test_device_conversion, devices=paddle_compatible_devices
754
+ )
755
+ add_function_test(TestPaddle, "test_from_paddle", test_from_paddle, devices=paddle_compatible_devices)
756
+ add_function_test(
757
+ TestPaddle, "test_from_paddle_slices", test_from_paddle_slices, devices=paddle_compatible_devices
758
+ )
759
+ add_function_test(
760
+ TestPaddle, "test_array_ctype_from_paddle", test_array_ctype_from_paddle, devices=paddle_compatible_devices
761
+ )
762
+ add_function_test(
763
+ TestPaddle,
764
+ "test_from_paddle_zero_strides",
765
+ test_from_paddle_zero_strides,
766
+ devices=paddle_compatible_devices,
767
+ )
768
+ add_function_test(TestPaddle, "test_to_paddle", test_to_paddle, devices=paddle_compatible_devices)
769
+ add_function_test(TestPaddle, "test_paddle_zerocopy", test_paddle_zerocopy, devices=paddle_compatible_devices)
770
+ add_function_test(TestPaddle, "test_paddle_autograd", test_paddle_autograd, devices=paddle_compatible_devices)
771
+ add_function_test(TestPaddle, "test_direct", test_direct, devices=paddle_compatible_devices)
772
+
773
+ # NOTE: Graph not supported now
774
+ # if paddle_compatible_cuda_devices:
775
+ # add_function_test(
776
+ # TestPaddle,
777
+ # "test_warp_graph_warp_stream",
778
+ # test_warp_graph_warp_stream,
779
+ # devices=paddle_compatible_cuda_devices,
780
+ # )
781
+ # add_function_test(
782
+ # TestPaddle,
783
+ # "test_warp_graph_paddle_stream",
784
+ # test_warp_graph_paddle_stream,
785
+ # devices=paddle_compatible_cuda_devices,
786
+ # )
787
+
788
+ # multi-GPU not supported yet.
789
+ # if len(paddle_compatible_cuda_devices) > 1:
790
+ # add_function_test(TestPaddle, "test_paddle_mgpu_from_paddle", test_paddle_mgpu_from_paddle)
791
+ # add_function_test(TestPaddle, "test_paddle_mgpu_to_paddle", test_paddle_mgpu_to_paddle)
792
+ # add_function_test(TestPaddle, "test_paddle_mgpu_interop", test_paddle_mgpu_interop)
793
+
794
+ except Exception as e:
795
+ print(f"Skipping Paddle tests due to exception: {e}")
796
+
797
+
798
+ if __name__ == "__main__":
799
+ wp.clear_kernel_cache()
800
+ unittest.main(verbosity=2)