warp-lang 1.0.2__py3-none-win_amd64.whl → 1.2.0__py3-none-win_amd64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of warp-lang might be problematic. Click here for more details.

Files changed (356) hide show
  1. warp/__init__.py +108 -97
  2. warp/__init__.pyi +1 -1
  3. warp/bin/warp-clang.dll +0 -0
  4. warp/bin/warp.dll +0 -0
  5. warp/build.py +88 -113
  6. warp/build_dll.py +383 -375
  7. warp/builtins.py +3693 -3354
  8. warp/codegen.py +2925 -2792
  9. warp/config.py +40 -36
  10. warp/constants.py +49 -45
  11. warp/context.py +5409 -5102
  12. warp/dlpack.py +442 -442
  13. warp/examples/__init__.py +16 -16
  14. warp/examples/assets/bear.usd +0 -0
  15. warp/examples/assets/bunny.usd +0 -0
  16. warp/examples/assets/cartpole.urdf +110 -110
  17. warp/examples/assets/crazyflie.usd +0 -0
  18. warp/examples/assets/cube.usd +0 -0
  19. warp/examples/assets/nv_ant.xml +92 -92
  20. warp/examples/assets/nv_humanoid.xml +183 -183
  21. warp/examples/assets/quadruped.urdf +267 -267
  22. warp/examples/assets/rocks.nvdb +0 -0
  23. warp/examples/assets/rocks.usd +0 -0
  24. warp/examples/assets/sphere.usd +0 -0
  25. warp/examples/benchmarks/benchmark_api.py +381 -383
  26. warp/examples/benchmarks/benchmark_cloth.py +278 -277
  27. warp/examples/benchmarks/benchmark_cloth_cupy.py +88 -88
  28. warp/examples/benchmarks/benchmark_cloth_jax.py +97 -100
  29. warp/examples/benchmarks/benchmark_cloth_numba.py +146 -142
  30. warp/examples/benchmarks/benchmark_cloth_numpy.py +77 -77
  31. warp/examples/benchmarks/benchmark_cloth_pytorch.py +86 -86
  32. warp/examples/benchmarks/benchmark_cloth_taichi.py +112 -112
  33. warp/examples/benchmarks/benchmark_cloth_warp.py +145 -146
  34. warp/examples/benchmarks/benchmark_launches.py +293 -295
  35. warp/examples/browse.py +29 -29
  36. warp/examples/core/example_dem.py +232 -219
  37. warp/examples/core/example_fluid.py +291 -267
  38. warp/examples/core/example_graph_capture.py +142 -126
  39. warp/examples/core/example_marching_cubes.py +186 -174
  40. warp/examples/core/example_mesh.py +172 -155
  41. warp/examples/core/example_mesh_intersect.py +203 -193
  42. warp/examples/core/example_nvdb.py +174 -170
  43. warp/examples/core/example_raycast.py +103 -90
  44. warp/examples/core/example_raymarch.py +197 -178
  45. warp/examples/core/example_render_opengl.py +183 -141
  46. warp/examples/core/example_sph.py +403 -387
  47. warp/examples/core/example_torch.py +219 -181
  48. warp/examples/core/example_wave.py +261 -248
  49. warp/examples/fem/bsr_utils.py +378 -380
  50. warp/examples/fem/example_apic_fluid.py +432 -389
  51. warp/examples/fem/example_burgers.py +262 -0
  52. warp/examples/fem/example_convection_diffusion.py +180 -168
  53. warp/examples/fem/example_convection_diffusion_dg.py +217 -209
  54. warp/examples/fem/example_deformed_geometry.py +175 -159
  55. warp/examples/fem/example_diffusion.py +199 -173
  56. warp/examples/fem/example_diffusion_3d.py +178 -152
  57. warp/examples/fem/example_diffusion_mgpu.py +219 -214
  58. warp/examples/fem/example_mixed_elasticity.py +242 -222
  59. warp/examples/fem/example_navier_stokes.py +257 -243
  60. warp/examples/fem/example_stokes.py +218 -192
  61. warp/examples/fem/example_stokes_transfer.py +263 -249
  62. warp/examples/fem/mesh_utils.py +133 -109
  63. warp/examples/fem/plot_utils.py +292 -287
  64. warp/examples/optim/example_bounce.py +258 -246
  65. warp/examples/optim/example_cloth_throw.py +220 -209
  66. warp/examples/optim/example_diffray.py +564 -536
  67. warp/examples/optim/example_drone.py +862 -835
  68. warp/examples/optim/example_inverse_kinematics.py +174 -168
  69. warp/examples/optim/example_inverse_kinematics_torch.py +183 -169
  70. warp/examples/optim/example_spring_cage.py +237 -231
  71. warp/examples/optim/example_trajectory.py +221 -199
  72. warp/examples/optim/example_walker.py +304 -293
  73. warp/examples/sim/example_cartpole.py +137 -129
  74. warp/examples/sim/example_cloth.py +194 -186
  75. warp/examples/sim/example_granular.py +122 -111
  76. warp/examples/sim/example_granular_collision_sdf.py +195 -186
  77. warp/examples/sim/example_jacobian_ik.py +234 -214
  78. warp/examples/sim/example_particle_chain.py +116 -105
  79. warp/examples/sim/example_quadruped.py +191 -180
  80. warp/examples/sim/example_rigid_chain.py +195 -187
  81. warp/examples/sim/example_rigid_contact.py +187 -177
  82. warp/examples/sim/example_rigid_force.py +125 -125
  83. warp/examples/sim/example_rigid_gyroscopic.py +107 -95
  84. warp/examples/sim/example_rigid_soft_contact.py +132 -122
  85. warp/examples/sim/example_soft_body.py +188 -177
  86. warp/fabric.py +337 -335
  87. warp/fem/__init__.py +61 -27
  88. warp/fem/cache.py +403 -388
  89. warp/fem/dirichlet.py +178 -179
  90. warp/fem/domain.py +262 -263
  91. warp/fem/field/__init__.py +100 -101
  92. warp/fem/field/field.py +148 -149
  93. warp/fem/field/nodal_field.py +298 -299
  94. warp/fem/field/restriction.py +22 -21
  95. warp/fem/field/test.py +180 -181
  96. warp/fem/field/trial.py +183 -183
  97. warp/fem/geometry/__init__.py +16 -19
  98. warp/fem/geometry/closest_point.py +69 -70
  99. warp/fem/geometry/deformed_geometry.py +270 -271
  100. warp/fem/geometry/element.py +748 -744
  101. warp/fem/geometry/geometry.py +184 -186
  102. warp/fem/geometry/grid_2d.py +380 -373
  103. warp/fem/geometry/grid_3d.py +437 -435
  104. warp/fem/geometry/hexmesh.py +953 -953
  105. warp/fem/geometry/nanogrid.py +455 -0
  106. warp/fem/geometry/partition.py +374 -376
  107. warp/fem/geometry/quadmesh_2d.py +532 -532
  108. warp/fem/geometry/tetmesh.py +840 -840
  109. warp/fem/geometry/trimesh_2d.py +577 -577
  110. warp/fem/integrate.py +1684 -1615
  111. warp/fem/operator.py +190 -191
  112. warp/fem/polynomial.py +214 -213
  113. warp/fem/quadrature/__init__.py +2 -2
  114. warp/fem/quadrature/pic_quadrature.py +243 -245
  115. warp/fem/quadrature/quadrature.py +295 -294
  116. warp/fem/space/__init__.py +179 -292
  117. warp/fem/space/basis_space.py +522 -489
  118. warp/fem/space/collocated_function_space.py +100 -105
  119. warp/fem/space/dof_mapper.py +236 -236
  120. warp/fem/space/function_space.py +148 -145
  121. warp/fem/space/grid_2d_function_space.py +148 -267
  122. warp/fem/space/grid_3d_function_space.py +167 -306
  123. warp/fem/space/hexmesh_function_space.py +253 -352
  124. warp/fem/space/nanogrid_function_space.py +202 -0
  125. warp/fem/space/partition.py +350 -350
  126. warp/fem/space/quadmesh_2d_function_space.py +261 -369
  127. warp/fem/space/restriction.py +161 -160
  128. warp/fem/space/shape/__init__.py +90 -15
  129. warp/fem/space/shape/cube_shape_function.py +728 -738
  130. warp/fem/space/shape/shape_function.py +102 -103
  131. warp/fem/space/shape/square_shape_function.py +611 -611
  132. warp/fem/space/shape/tet_shape_function.py +565 -567
  133. warp/fem/space/shape/triangle_shape_function.py +429 -429
  134. warp/fem/space/tetmesh_function_space.py +224 -292
  135. warp/fem/space/topology.py +297 -295
  136. warp/fem/space/trimesh_2d_function_space.py +153 -221
  137. warp/fem/types.py +77 -77
  138. warp/fem/utils.py +495 -495
  139. warp/jax.py +166 -141
  140. warp/jax_experimental.py +341 -339
  141. warp/native/array.h +1081 -1025
  142. warp/native/builtin.h +1603 -1560
  143. warp/native/bvh.cpp +402 -398
  144. warp/native/bvh.cu +533 -525
  145. warp/native/bvh.h +430 -429
  146. warp/native/clang/clang.cpp +496 -464
  147. warp/native/crt.cpp +42 -32
  148. warp/native/crt.h +352 -335
  149. warp/native/cuda_crt.h +1049 -1049
  150. warp/native/cuda_util.cpp +549 -540
  151. warp/native/cuda_util.h +288 -203
  152. warp/native/cutlass_gemm.cpp +34 -34
  153. warp/native/cutlass_gemm.cu +372 -372
  154. warp/native/error.cpp +66 -66
  155. warp/native/error.h +27 -27
  156. warp/native/exports.h +187 -0
  157. warp/native/fabric.h +228 -228
  158. warp/native/hashgrid.cpp +301 -278
  159. warp/native/hashgrid.cu +78 -77
  160. warp/native/hashgrid.h +227 -227
  161. warp/native/initializer_array.h +32 -32
  162. warp/native/intersect.h +1204 -1204
  163. warp/native/intersect_adj.h +365 -365
  164. warp/native/intersect_tri.h +322 -322
  165. warp/native/marching.cpp +2 -2
  166. warp/native/marching.cu +497 -497
  167. warp/native/marching.h +2 -2
  168. warp/native/mat.h +1545 -1498
  169. warp/native/matnn.h +333 -333
  170. warp/native/mesh.cpp +203 -203
  171. warp/native/mesh.cu +292 -293
  172. warp/native/mesh.h +1887 -1887
  173. warp/native/nanovdb/GridHandle.h +366 -0
  174. warp/native/nanovdb/HostBuffer.h +590 -0
  175. warp/native/nanovdb/NanoVDB.h +6624 -4782
  176. warp/native/nanovdb/PNanoVDB.h +3390 -2553
  177. warp/native/noise.h +850 -850
  178. warp/native/quat.h +1112 -1085
  179. warp/native/rand.h +303 -299
  180. warp/native/range.h +108 -108
  181. warp/native/reduce.cpp +156 -156
  182. warp/native/reduce.cu +348 -348
  183. warp/native/runlength_encode.cpp +61 -61
  184. warp/native/runlength_encode.cu +46 -46
  185. warp/native/scan.cpp +30 -30
  186. warp/native/scan.cu +36 -36
  187. warp/native/scan.h +7 -7
  188. warp/native/solid_angle.h +442 -442
  189. warp/native/sort.cpp +94 -94
  190. warp/native/sort.cu +97 -97
  191. warp/native/sort.h +14 -14
  192. warp/native/sparse.cpp +337 -337
  193. warp/native/sparse.cu +544 -544
  194. warp/native/spatial.h +630 -630
  195. warp/native/svd.h +562 -562
  196. warp/native/temp_buffer.h +30 -30
  197. warp/native/vec.h +1177 -1133
  198. warp/native/volume.cpp +529 -297
  199. warp/native/volume.cu +58 -32
  200. warp/native/volume.h +960 -538
  201. warp/native/volume_builder.cu +446 -425
  202. warp/native/volume_builder.h +34 -19
  203. warp/native/volume_impl.h +61 -0
  204. warp/native/warp.cpp +1057 -1052
  205. warp/native/warp.cu +2949 -2828
  206. warp/native/warp.h +321 -305
  207. warp/optim/__init__.py +9 -9
  208. warp/optim/adam.py +120 -120
  209. warp/optim/linear.py +1104 -939
  210. warp/optim/sgd.py +104 -92
  211. warp/render/__init__.py +10 -10
  212. warp/render/render_opengl.py +3356 -3204
  213. warp/render/render_usd.py +768 -749
  214. warp/render/utils.py +152 -150
  215. warp/sim/__init__.py +52 -59
  216. warp/sim/articulation.py +685 -685
  217. warp/sim/collide.py +1594 -1590
  218. warp/sim/import_mjcf.py +489 -481
  219. warp/sim/import_snu.py +220 -221
  220. warp/sim/import_urdf.py +536 -516
  221. warp/sim/import_usd.py +887 -881
  222. warp/sim/inertia.py +316 -317
  223. warp/sim/integrator.py +234 -233
  224. warp/sim/integrator_euler.py +1956 -1956
  225. warp/sim/integrator_featherstone.py +1917 -1991
  226. warp/sim/integrator_xpbd.py +3288 -3312
  227. warp/sim/model.py +4473 -4314
  228. warp/sim/particles.py +113 -112
  229. warp/sim/render.py +417 -403
  230. warp/sim/utils.py +413 -410
  231. warp/sparse.py +1289 -1227
  232. warp/stubs.py +2192 -2469
  233. warp/tape.py +1162 -225
  234. warp/tests/__init__.py +1 -1
  235. warp/tests/__main__.py +4 -4
  236. warp/tests/assets/test_index_grid.nvdb +0 -0
  237. warp/tests/assets/torus.usda +105 -105
  238. warp/tests/aux_test_class_kernel.py +26 -26
  239. warp/tests/aux_test_compile_consts_dummy.py +10 -10
  240. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -21
  241. warp/tests/aux_test_dependent.py +20 -22
  242. warp/tests/aux_test_grad_customs.py +21 -23
  243. warp/tests/aux_test_reference.py +9 -11
  244. warp/tests/aux_test_reference_reference.py +8 -10
  245. warp/tests/aux_test_square.py +15 -17
  246. warp/tests/aux_test_unresolved_func.py +14 -14
  247. warp/tests/aux_test_unresolved_symbol.py +14 -14
  248. warp/tests/disabled_kinematics.py +237 -239
  249. warp/tests/run_coverage_serial.py +31 -31
  250. warp/tests/test_adam.py +155 -157
  251. warp/tests/test_arithmetic.py +1088 -1124
  252. warp/tests/test_array.py +2415 -2326
  253. warp/tests/test_array_reduce.py +148 -150
  254. warp/tests/test_async.py +666 -656
  255. warp/tests/test_atomic.py +139 -141
  256. warp/tests/test_bool.py +212 -149
  257. warp/tests/test_builtins_resolution.py +1290 -1292
  258. warp/tests/test_bvh.py +162 -171
  259. warp/tests/test_closest_point_edge_edge.py +227 -228
  260. warp/tests/test_codegen.py +562 -553
  261. warp/tests/test_compile_consts.py +217 -101
  262. warp/tests/test_conditional.py +244 -246
  263. warp/tests/test_copy.py +230 -215
  264. warp/tests/test_ctypes.py +630 -632
  265. warp/tests/test_dense.py +65 -67
  266. warp/tests/test_devices.py +89 -98
  267. warp/tests/test_dlpack.py +528 -529
  268. warp/tests/test_examples.py +403 -378
  269. warp/tests/test_fabricarray.py +952 -955
  270. warp/tests/test_fast_math.py +60 -54
  271. warp/tests/test_fem.py +1298 -1278
  272. warp/tests/test_fp16.py +128 -130
  273. warp/tests/test_func.py +336 -337
  274. warp/tests/test_generics.py +596 -571
  275. warp/tests/test_grad.py +885 -640
  276. warp/tests/test_grad_customs.py +331 -336
  277. warp/tests/test_hash_grid.py +208 -164
  278. warp/tests/test_import.py +37 -39
  279. warp/tests/test_indexedarray.py +1132 -1134
  280. warp/tests/test_intersect.py +65 -67
  281. warp/tests/test_jax.py +305 -307
  282. warp/tests/test_large.py +169 -164
  283. warp/tests/test_launch.py +352 -354
  284. warp/tests/test_lerp.py +217 -261
  285. warp/tests/test_linear_solvers.py +189 -171
  286. warp/tests/test_lvalue.py +419 -493
  287. warp/tests/test_marching_cubes.py +63 -65
  288. warp/tests/test_mat.py +1799 -1827
  289. warp/tests/test_mat_lite.py +113 -115
  290. warp/tests/test_mat_scalar_ops.py +2905 -2889
  291. warp/tests/test_math.py +124 -193
  292. warp/tests/test_matmul.py +498 -499
  293. warp/tests/test_matmul_lite.py +408 -410
  294. warp/tests/test_mempool.py +186 -190
  295. warp/tests/test_mesh.py +281 -324
  296. warp/tests/test_mesh_query_aabb.py +226 -241
  297. warp/tests/test_mesh_query_point.py +690 -702
  298. warp/tests/test_mesh_query_ray.py +290 -303
  299. warp/tests/test_mlp.py +274 -276
  300. warp/tests/test_model.py +108 -110
  301. warp/tests/test_module_hashing.py +111 -0
  302. warp/tests/test_modules_lite.py +36 -39
  303. warp/tests/test_multigpu.py +161 -163
  304. warp/tests/test_noise.py +244 -248
  305. warp/tests/test_operators.py +248 -250
  306. warp/tests/test_options.py +121 -125
  307. warp/tests/test_peer.py +131 -137
  308. warp/tests/test_pinned.py +76 -78
  309. warp/tests/test_print.py +52 -54
  310. warp/tests/test_quat.py +2084 -2086
  311. warp/tests/test_rand.py +324 -288
  312. warp/tests/test_reload.py +207 -217
  313. warp/tests/test_rounding.py +177 -179
  314. warp/tests/test_runlength_encode.py +188 -190
  315. warp/tests/test_sim_grad.py +241 -0
  316. warp/tests/test_sim_kinematics.py +89 -97
  317. warp/tests/test_smoothstep.py +166 -168
  318. warp/tests/test_snippet.py +303 -266
  319. warp/tests/test_sparse.py +466 -460
  320. warp/tests/test_spatial.py +2146 -2148
  321. warp/tests/test_special_values.py +362 -0
  322. warp/tests/test_streams.py +484 -473
  323. warp/tests/test_struct.py +708 -675
  324. warp/tests/test_tape.py +171 -148
  325. warp/tests/test_torch.py +741 -743
  326. warp/tests/test_transient_module.py +85 -87
  327. warp/tests/test_types.py +554 -659
  328. warp/tests/test_utils.py +488 -499
  329. warp/tests/test_vec.py +1262 -1268
  330. warp/tests/test_vec_lite.py +71 -73
  331. warp/tests/test_vec_scalar_ops.py +2097 -2099
  332. warp/tests/test_verify_fp.py +92 -94
  333. warp/tests/test_volume.py +961 -736
  334. warp/tests/test_volume_write.py +338 -265
  335. warp/tests/unittest_serial.py +38 -37
  336. warp/tests/unittest_suites.py +367 -359
  337. warp/tests/unittest_utils.py +434 -578
  338. warp/tests/unused_test_misc.py +69 -71
  339. warp/tests/walkthrough_debug.py +85 -85
  340. warp/thirdparty/appdirs.py +598 -598
  341. warp/thirdparty/dlpack.py +143 -143
  342. warp/thirdparty/unittest_parallel.py +563 -561
  343. warp/torch.py +321 -295
  344. warp/types.py +4941 -4450
  345. warp/utils.py +1008 -821
  346. {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/LICENSE.md +126 -126
  347. {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/METADATA +365 -400
  348. warp_lang-1.2.0.dist-info/RECORD +359 -0
  349. warp/examples/assets/cube.usda +0 -42
  350. warp/examples/assets/sphere.usda +0 -56
  351. warp/examples/assets/torus.usda +0 -105
  352. warp/examples/fem/example_convection_diffusion_dg0.py +0 -194
  353. warp/native/nanovdb/PNanoVDBWrite.h +0 -295
  354. warp_lang-1.0.2.dist-info/RECORD +0 -352
  355. {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/WHEEL +0 -0
  356. {warp_lang-1.0.2.dist-info → warp_lang-1.2.0.dist-info}/top_level.txt +0 -0
@@ -1,955 +1,952 @@
1
- # Copyright (c) 2023 NVIDIA CORPORATION. All rights reserved.
2
- # NVIDIA CORPORATION and its licensors retain all intellectual property
3
- # and proprietary rights in and to this software, related documentation
4
- # and any modifications thereto. Any use, reproduction, disclosure or
5
- # distribution of this software and related documentation without an express
6
- # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
-
8
- import math
9
- import unittest
10
- from typing import Any
11
-
12
- import numpy as np
13
-
14
- import warp as wp
15
- from warp.tests.unittest_utils import *
16
-
17
- wp.init()
18
-
19
-
20
- # types to test fabric arrays
21
- _fabric_types = [
22
- *wp.types.scalar_types,
23
- *[wp.types.vector(2, T) for T in wp.types.scalar_types],
24
- *[wp.types.vector(3, T) for T in wp.types.scalar_types],
25
- *[wp.types.vector(4, T) for T in wp.types.scalar_types],
26
- *[wp.types.matrix((2, 2), T) for T in wp.types.scalar_types],
27
- *[wp.types.matrix((3, 3), T) for T in wp.types.scalar_types],
28
- *[wp.types.matrix((4, 4), T) for T in wp.types.scalar_types],
29
- *[wp.types.quaternion(T) for T in wp.types.float_types],
30
- ]
31
-
32
-
33
- def _warp_type_to_fabric(dtype, is_array=False):
34
- scalar_map = {
35
- wp.bool: "b",
36
- wp.int8: "i1",
37
- wp.int16: "i2",
38
- wp.int32: "i4",
39
- wp.int64: "i8",
40
- wp.uint8: "u1",
41
- wp.uint16: "u2",
42
- wp.uint32: "u4",
43
- wp.uint64: "u8",
44
- wp.float16: "f2",
45
- wp.float32: "f4",
46
- wp.float64: "f8",
47
- }
48
-
49
- if hasattr(dtype, "_wp_scalar_type_"):
50
- type_str = scalar_map[dtype._wp_scalar_type_]
51
- if len(dtype._shape_) == 1:
52
- role = "vector"
53
- else:
54
- role = "matrix"
55
- else:
56
- type_str = scalar_map[dtype]
57
- role = ""
58
-
59
- if is_array:
60
- array_depth = 1
61
- else:
62
- array_depth = 0
63
-
64
- return (True, type_str, dtype._length_, array_depth, role)
65
-
66
-
67
- # returns a fabric array interface constructed from a regular array
68
- def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list = None, copy=False):
69
- assert isinstance(data, wp.array)
70
- assert data.ndim == 1
71
-
72
- assert isinstance(attrib, str)
73
-
74
- if copy:
75
- data = wp.clone(data)
76
-
77
- if bucket_sizes is not None:
78
- assert hasattr(bucket_sizes, "__len__")
79
-
80
- # verify total size
81
- total_size = 0
82
- for bucket_size in bucket_sizes:
83
- total_size += bucket_size
84
-
85
- if total_size != data.size:
86
- raise RuntimeError("Bucket sizes don't add up to the size of data array")
87
-
88
- elif data.size > 0:
89
- rng = np.random.default_rng(123)
90
-
91
- # generate random bucket sizes
92
- bucket_min = 1
93
- bucket_max = math.ceil(0.5 * data.size)
94
- total_size = data.size
95
- size_remaining = total_size
96
-
97
- bucket_sizes = []
98
- while size_remaining >= bucket_max:
99
- bucket_size = rng.integers(bucket_min, high=bucket_max, dtype=int)
100
- bucket_sizes.append(bucket_size)
101
- size_remaining -= bucket_size
102
-
103
- if size_remaining > 0:
104
- bucket_sizes.append(size_remaining)
105
-
106
- else:
107
- # empty data array
108
- bucket_sizes = []
109
-
110
- dtype_size = wp.types.type_size_in_bytes(data.dtype)
111
- p = int(data.ptr) if data.ptr else 0
112
- pointers = []
113
- counts = []
114
- for bucket_size in bucket_sizes:
115
- pointers.append(p)
116
- counts.append(bucket_size)
117
- p += bucket_size * dtype_size
118
-
119
- attrib_info = {}
120
-
121
- attrib_info["type"] = _warp_type_to_fabric(data.dtype)
122
- attrib_info["access"] = 2 # ReadWrite
123
- attrib_info["pointers"] = pointers
124
- attrib_info["counts"] = counts
125
-
126
- iface = {}
127
- iface["version"] = 1
128
- iface["device"] = str(data.device)
129
- iface["attribs"] = {attrib: attrib_info}
130
- iface["_ref"] = data # backref to keep the array alive
131
-
132
- return iface
133
-
134
-
135
- # returns a fabric array array interface constructed from a list of regular arrays
136
- def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list = None):
137
- # data should be a list of arrays
138
- assert isinstance(data, list)
139
-
140
- num_arrays = len(data)
141
- assert num_arrays > 0
142
-
143
- device = data[0].device
144
- dtype = data[0].dtype
145
-
146
- assert isinstance(attrib, str)
147
-
148
- if bucket_sizes is not None:
149
- assert hasattr(bucket_sizes, "__len__")
150
-
151
- # verify total size
152
- total_size = 0
153
- for bucket_size in bucket_sizes:
154
- total_size += bucket_size
155
-
156
- if total_size != num_arrays:
157
- raise RuntimeError("Bucket sizes don't add up to the number of given arrays")
158
-
159
- else:
160
- rng = np.random.default_rng(123)
161
-
162
- # generate random bucket sizes
163
- bucket_min = 1
164
- bucket_max = math.ceil(0.5 * num_arrays)
165
- total_size = num_arrays
166
- size_remaining = total_size
167
-
168
- bucket_sizes = []
169
- while size_remaining >= bucket_max:
170
- bucket_size = rng.integers(bucket_min, high=bucket_max, dtype=int)
171
- bucket_sizes.append(bucket_size)
172
- size_remaining -= bucket_size
173
-
174
- if size_remaining > 0:
175
- bucket_sizes.append(size_remaining)
176
-
177
- # initialize array of pointers to arrays and their lengths
178
- _array_pointers = []
179
- _array_lengths = []
180
- for i in range(num_arrays):
181
- _array_pointers.append(data[i].ptr)
182
- _array_lengths.append(data[i].size)
183
-
184
- array_pointers = wp.array(_array_pointers, dtype=wp.uint64, device=device)
185
- pointer_size = wp.types.type_size_in_bytes(array_pointers.dtype)
186
-
187
- lengths = wp.array(_array_lengths, dtype=wp.uint64, device=device)
188
- length_size = wp.types.type_size_in_bytes(lengths.dtype)
189
-
190
- p_pointers = int(array_pointers.ptr)
191
- p_lengths = int(lengths.ptr)
192
- pointers = []
193
- counts = []
194
- array_lengths = []
195
- for bucket_size in bucket_sizes:
196
- pointers.append(p_pointers)
197
- counts.append(bucket_size)
198
- array_lengths.append(p_lengths)
199
- p_pointers += bucket_size * pointer_size
200
- p_lengths += bucket_size * length_size
201
-
202
- attrib_info = {}
203
-
204
- attrib_info["type"] = _warp_type_to_fabric(dtype, is_array=True)
205
- attrib_info["access"] = 2 # ReadWrite
206
- attrib_info["pointers"] = pointers
207
- attrib_info["counts"] = counts
208
- attrib_info["array_lengths"] = array_lengths
209
-
210
- iface = {}
211
- iface["version"] = 1
212
- iface["device"] = str(device)
213
- iface["attribs"] = {attrib: attrib_info}
214
- iface["_ref"] = data # backref to keep the data arrays alive
215
- iface["_ref_pointers"] = array_pointers # backref to keep the array pointers alive
216
- iface["_ref_lengths"] = lengths # backref to keep the lengths array alive
217
-
218
- return iface
219
-
220
-
221
- @wp.kernel
222
- def fa_kernel(a: wp.fabricarray(dtype=float), expected: wp.array(dtype=float)):
223
- i = wp.tid()
224
-
225
- wp.expect_eq(a[i], expected[i])
226
-
227
- a[i] = 2.0 * a[i]
228
-
229
- wp.atomic_add(a, i, 1.0)
230
-
231
- wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)
232
-
233
-
234
- @wp.kernel
235
- def fa_kernel_indexed(a: wp.indexedfabricarray(dtype=float), expected: wp.indexedarray(dtype=float)):
236
- i = wp.tid()
237
-
238
- wp.expect_eq(a[i], expected[i])
239
-
240
- a[i] = 2.0 * a[i]
241
-
242
- wp.atomic_add(a, i, 1.0)
243
-
244
- wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)
245
-
246
-
247
- def test_fabricarray_kernel(test, device):
248
- data = wp.array(data=np.arange(100, dtype=np.float32), device=device)
249
- iface = _create_fabric_array_interface(data, "foo", copy=True)
250
- fa = wp.fabricarray(data=iface, attrib="foo")
251
-
252
- test.assertEqual(fa.dtype, data.dtype)
253
- test.assertEqual(fa.ndim, 1)
254
- test.assertEqual(fa.shape, data.shape)
255
- test.assertEqual(fa.size, data.size)
256
-
257
- wp.launch(fa_kernel, dim=fa.size, inputs=[fa, data], device=device)
258
-
259
- # reset data
260
- wp.copy(fa, data)
261
-
262
- # test indexed
263
- indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
264
- ifa = fa[indices]
265
- idata = data[indices]
266
-
267
- test.assertEqual(ifa.dtype, idata.dtype)
268
- test.assertEqual(ifa.ndim, 1)
269
- test.assertEqual(ifa.shape, idata.shape)
270
- test.assertEqual(ifa.size, idata.size)
271
-
272
- wp.launch(fa_kernel_indexed, dim=ifa.size, inputs=[ifa, idata], device=device)
273
-
274
- wp.synchronize_device(device)
275
-
276
-
277
- @wp.kernel
278
- def fa_generic_dtype_kernel(a: wp.fabricarray(dtype=Any), b: wp.fabricarray(dtype=Any)):
279
- i = wp.tid()
280
- b[i] = a[i] + a[i]
281
-
282
-
283
- @wp.kernel
284
- def fa_generic_dtype_kernel_indexed(a: wp.indexedfabricarray(dtype=Any), b: wp.indexedfabricarray(dtype=Any)):
285
- i = wp.tid()
286
- b[i] = a[i] + a[i]
287
-
288
-
289
- def test_fabricarray_generic_dtype(test, device):
290
- for T in _fabric_types:
291
- if hasattr(T, "_wp_scalar_type_"):
292
- nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
293
- else:
294
- nptype = wp.types.warp_type_to_np_dtype[T]
295
-
296
- data = wp.array(data=np.arange(10, dtype=nptype), device=device)
297
- data_iface = _create_fabric_array_interface(data, "foo", copy=True)
298
- fa = wp.fabricarray(data=data_iface, attrib="foo")
299
-
300
- result = wp.zeros_like(data)
301
- result_iface = _create_fabric_array_interface(result, "foo", copy=True)
302
- fb = wp.fabricarray(data=result_iface, attrib="foo")
303
-
304
- test.assertEqual(fa.dtype, fb.dtype)
305
- test.assertEqual(fa.ndim, fb.ndim)
306
- test.assertEqual(fa.shape, fb.shape)
307
- test.assertEqual(fa.size, fb.size)
308
-
309
- wp.launch(fa_generic_dtype_kernel, dim=fa.size, inputs=[fa, fb], device=device)
310
-
311
- assert_np_equal(fb.numpy(), 2 * fa.numpy())
312
-
313
- # reset data
314
- wp.copy(fa, data)
315
- wp.copy(fb, result)
316
-
317
- # test indexed
318
- indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
319
- ifa = fa[indices]
320
- ifb = fb[indices]
321
-
322
- test.assertEqual(ifa.dtype, ifb.dtype)
323
- test.assertEqual(ifa.ndim, ifb.ndim)
324
- test.assertEqual(ifa.shape, ifb.shape)
325
- test.assertEqual(ifa.size, ifb.size)
326
-
327
- wp.launch(fa_generic_dtype_kernel_indexed, dim=ifa.size, inputs=[ifa, ifb], device=device)
328
-
329
- assert_np_equal(ifb.numpy(), 2 * ifa.numpy())
330
-
331
-
332
- @wp.kernel
333
- def fa_generic_array_kernel(a: Any, b: Any):
334
- i = wp.tid()
335
- b[i] = a[i] + a[i]
336
-
337
-
338
- def test_fabricarray_generic_array(test, device):
339
- for T in _fabric_types:
340
- if hasattr(T, "_wp_scalar_type_"):
341
- nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
342
- else:
343
- nptype = wp.types.warp_type_to_np_dtype[T]
344
-
345
- data = wp.array(data=np.arange(100, dtype=nptype), device=device)
346
- data_iface = _create_fabric_array_interface(data, "foo", copy=True)
347
- fa = wp.fabricarray(data=data_iface, attrib="foo")
348
-
349
- result = wp.zeros_like(data)
350
- result_iface = _create_fabric_array_interface(result, "foo", copy=True)
351
- fb = wp.fabricarray(data=result_iface, attrib="foo")
352
-
353
- test.assertEqual(fa.dtype, fb.dtype)
354
- test.assertEqual(fa.ndim, fb.ndim)
355
- test.assertEqual(fa.shape, fb.shape)
356
- test.assertEqual(fa.size, fb.size)
357
-
358
- wp.launch(fa_generic_array_kernel, dim=fa.size, inputs=[fa, fb], device=device)
359
-
360
- assert_np_equal(fb.numpy(), 2 * fa.numpy())
361
-
362
- # reset data
363
- wp.copy(fa, data)
364
- wp.copy(fb, result)
365
-
366
- # test indexed
367
- indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
368
- ifa = fa[indices]
369
- ifb = fb[indices]
370
-
371
- test.assertEqual(ifa.dtype, ifb.dtype)
372
- test.assertEqual(ifa.ndim, ifb.ndim)
373
- test.assertEqual(ifa.shape, ifb.shape)
374
- test.assertEqual(ifa.size, ifb.size)
375
-
376
- wp.launch(fa_generic_array_kernel, dim=ifa.size, inputs=[ifa, ifb], device=device)
377
-
378
- assert_np_equal(ifb.numpy(), 2 * ifa.numpy())
379
-
380
-
381
- def test_fabricarray_empty(test, device):
382
- # Test whether common operations work with empty (zero-sized) indexed arrays
383
- # without throwing exceptions.
384
-
385
- def test_empty_ops(nrows, ncols, wptype, nptype):
386
- # scalar, vector, or matrix
387
- if ncols > 0:
388
- if nrows > 0:
389
- wptype = wp.types.matrix((nrows, ncols), wptype)
390
- else:
391
- wptype = wp.types.vector(ncols, wptype)
392
- dtype_shape = wptype._shape_
393
- else:
394
- dtype_shape = ()
395
-
396
- fill_value = wptype(42)
397
-
398
- # create an empty data array
399
- data = wp.empty(0, dtype=wptype, device=device)
400
- iface = _create_fabric_array_interface(data, "foo", copy=True)
401
- fa = wp.fabricarray(data=iface, attrib="foo")
402
-
403
- test.assertEqual(fa.size, 0)
404
- test.assertEqual(fa.shape, (0,))
405
-
406
- # all of these methods should succeed with zero-sized arrays
407
- fa.zero_()
408
- fa.fill_(fill_value)
409
- fb = fa.contiguous()
410
-
411
- fb = wp.empty_like(fa)
412
- fb = wp.zeros_like(fa)
413
- fb = wp.full_like(fa, fill_value)
414
- fb = wp.clone(fa)
415
-
416
- wp.copy(fa, fb)
417
- fa.assign(fb)
418
-
419
- na = fa.numpy()
420
- test.assertEqual(na.size, 0)
421
- test.assertEqual(na.shape, (0, *dtype_shape))
422
- test.assertEqual(na.dtype, nptype)
423
-
424
- test.assertEqual(fa.list(), [])
425
-
426
- # test indexed
427
-
428
- # create a zero-sized array of indices
429
- indices = wp.empty(0, dtype=int, device=device)
430
-
431
- ifa = fa[indices]
432
-
433
- test.assertEqual(ifa.size, 0)
434
- test.assertEqual(ifa.shape, (0,))
435
-
436
- # all of these methods should succeed with zero-sized arrays
437
- ifa.zero_()
438
- ifa.fill_(fill_value)
439
- ifb = ifa.contiguous()
440
-
441
- ifb = wp.empty_like(ifa)
442
- ifb = wp.zeros_like(ifa)
443
- ifb = wp.full_like(ifa, fill_value)
444
- ifb = wp.clone(ifa)
445
-
446
- wp.copy(ifa, ifb)
447
- ifa.assign(ifb)
448
-
449
- na = ifa.numpy()
450
- test.assertEqual(na.size, 0)
451
- test.assertEqual(na.shape, (0, *dtype_shape))
452
- test.assertEqual(na.dtype, nptype)
453
-
454
- test.assertEqual(ifa.list(), [])
455
-
456
- # test with scalars, vectors, and matrices
457
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
458
- # scalars
459
- test_empty_ops(0, 0, wptype, nptype)
460
-
461
- for ncols in [2, 3, 4, 5]:
462
- # vectors
463
- test_empty_ops(0, ncols, wptype, nptype)
464
- # square matrices (the Fabric interface only supports square matrices right now)
465
- test_empty_ops(ncols, ncols, wptype, nptype)
466
-
467
-
468
- def test_fabricarray_fill_scalar(test, device):
469
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
470
- # create a data array
471
- data = wp.zeros(100, dtype=wptype, device=device)
472
- iface = _create_fabric_array_interface(data, "foo", copy=True)
473
- fa = wp.fabricarray(data=iface, attrib="foo")
474
-
475
- assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))
476
-
477
- # fill with int value
478
- fill_value = 42
479
- fa.fill_(fill_value)
480
- assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value, dtype=nptype))
481
-
482
- fa.zero_()
483
- assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))
484
-
485
- if wptype in wp.types.float_types:
486
- # fill with float value
487
- fill_value = 13.37
488
- fa.fill_(fill_value)
489
- assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value, dtype=nptype))
490
-
491
- # fill with Warp scalar value
492
- fill_value = wptype(17)
493
- fa.fill_(fill_value)
494
- assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value.value, dtype=nptype))
495
-
496
- # reset data
497
- wp.copy(fa, data)
498
-
499
- # test indexed
500
- indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
501
- ifa = fa[indices1]
502
-
503
- # ensure that the other indices remain unchanged
504
- indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
505
- ifb = fa[indices2]
506
-
507
- assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
508
- assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
509
-
510
- # fill with int value
511
- fill_value = 42
512
- ifa.fill_(fill_value)
513
- assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value, dtype=nptype))
514
- assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
515
-
516
- ifa.zero_()
517
- assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
518
- assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
519
-
520
- if wptype in wp.types.float_types:
521
- # fill with float value
522
- fill_value = 13.37
523
- ifa.fill_(fill_value)
524
- assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value, dtype=nptype))
525
- assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
526
-
527
- # fill with Warp scalar value
528
- fill_value = wptype(17)
529
- ifa.fill_(fill_value)
530
- assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value.value, dtype=nptype))
531
- assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
532
-
533
-
534
- def test_fabricarray_fill_vector(test, device):
535
- # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
536
-
537
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
538
- # vector types
539
- vector_types = [
540
- wp.types.vector(2, wptype),
541
- wp.types.vector(3, wptype),
542
- wp.types.vector(4, wptype),
543
- wp.types.vector(5, wptype),
544
- ]
545
-
546
- for vec_type in vector_types:
547
- vec_len = vec_type._length_
548
-
549
- data = wp.zeros(100, dtype=vec_type, device=device)
550
- iface = _create_fabric_array_interface(data, "foo", copy=True)
551
- fa = wp.fabricarray(data=iface, attrib="foo")
552
-
553
- assert_np_equal(fa.numpy(), np.zeros((*fa.shape, vec_len), dtype=nptype))
554
-
555
- # fill with int scalar
556
- fill_value = 42
557
- fa.fill_(fill_value)
558
- assert_np_equal(fa.numpy(), np.full((*fa.shape, vec_len), fill_value, dtype=nptype))
559
-
560
- # test zeroing
561
- fa.zero_()
562
- assert_np_equal(fa.numpy(), np.zeros((*fa.shape, vec_len), dtype=nptype))
563
-
564
- # vector values can be passed as a list, numpy array, or Warp vector instance
565
- fill_list = [17, 42, 99, 101, 127][:vec_len]
566
- fill_arr = np.array(fill_list, dtype=nptype)
567
- fill_vec = vec_type(fill_list)
568
-
569
- expected = np.tile(fill_arr, fa.size).reshape((*fa.shape, vec_len))
570
-
571
- # fill with list of vector length
572
- fa.fill_(fill_list)
573
- assert_np_equal(fa.numpy(), expected)
574
-
575
- # clear
576
- fa.zero_()
577
-
578
- # fill with numpy array of vector length
579
- fa.fill_(fill_arr)
580
- assert_np_equal(fa.numpy(), expected)
581
-
582
- # clear
583
- fa.zero_()
584
-
585
- # fill with vec instance
586
- fa.fill_(fill_vec)
587
- assert_np_equal(fa.numpy(), expected)
588
-
589
- if wptype in wp.types.float_types:
590
- # fill with float scalar
591
- fill_value = 13.37
592
- fa.fill_(fill_value)
593
- assert_np_equal(fa.numpy(), np.full((*fa.shape, vec_len), fill_value, dtype=nptype))
594
-
595
- # fill with float list of vector length
596
- fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
597
-
598
- fa.fill_(fill_list)
599
-
600
- expected = np.tile(np.array(fill_list, dtype=nptype), fa.size).reshape((*fa.shape, vec_len))
601
-
602
- assert_np_equal(fa.numpy(), expected)
603
-
604
- # reset data
605
- wp.copy(fa, data)
606
-
607
- # test indexed
608
- indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
609
- ifa = fa[indices1]
610
-
611
- # ensure that the other indices remain unchanged
612
- indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
613
- ifb = fa[indices2]
614
-
615
- assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, vec_len), dtype=nptype))
616
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
617
-
618
- # fill with int scalar
619
- fill_value = 42
620
- ifa.fill_(fill_value)
621
- assert_np_equal(ifa.numpy(), np.full((*ifa.shape, vec_len), fill_value, dtype=nptype))
622
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
623
-
624
- # test zeroing
625
- ifa.zero_()
626
- assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, vec_len), dtype=nptype))
627
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
628
-
629
- # vector values can be passed as a list, numpy array, or Warp vector instance
630
- fill_list = [17, 42, 99, 101, 127][:vec_len]
631
- fill_arr = np.array(fill_list, dtype=nptype)
632
- fill_vec = vec_type(fill_list)
633
-
634
- expected = np.tile(fill_arr, ifa.size).reshape((*ifa.shape, vec_len))
635
-
636
- # fill with list of vector length
637
- ifa.fill_(fill_list)
638
- assert_np_equal(ifa.numpy(), expected)
639
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
640
-
641
- # clear
642
- ifa.zero_()
643
-
644
- # fill with numpy array of vector length
645
- ifa.fill_(fill_arr)
646
- assert_np_equal(ifa.numpy(), expected)
647
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
648
-
649
- # clear
650
- ifa.zero_()
651
-
652
- # fill with vec instance
653
- ifa.fill_(fill_vec)
654
- assert_np_equal(ifa.numpy(), expected)
655
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
656
-
657
- if wptype in wp.types.float_types:
658
- # fill with float scalar
659
- fill_value = 13.37
660
- ifa.fill_(fill_value)
661
- assert_np_equal(ifa.numpy(), np.full((*ifa.shape, vec_len), fill_value, dtype=nptype))
662
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
663
-
664
- # fill with float list of vector length
665
- fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
666
-
667
- ifa.fill_(fill_list)
668
-
669
- expected = np.tile(np.array(fill_list, dtype=nptype), ifa.size).reshape((*ifa.shape, vec_len))
670
-
671
- assert_np_equal(ifa.numpy(), expected)
672
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
673
-
674
-
675
- def test_fabricarray_fill_matrix(test, device):
676
- # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
677
-
678
- for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
679
- # matrix types
680
- matrix_types = [
681
- # square matrices only
682
- wp.types.matrix((2, 2), wptype),
683
- wp.types.matrix((3, 3), wptype),
684
- wp.types.matrix((4, 4), wptype),
685
- wp.types.matrix((5, 5), wptype),
686
- ]
687
-
688
- for mat_type in matrix_types:
689
- mat_len = mat_type._length_
690
- mat_shape = mat_type._shape_
691
-
692
- data = wp.zeros(100, dtype=mat_type, device=device)
693
- iface = _create_fabric_array_interface(data, "foo", copy=True)
694
- fa = wp.fabricarray(data=iface, attrib="foo")
695
-
696
- assert_np_equal(fa.numpy(), np.zeros((*fa.shape, *mat_shape), dtype=nptype))
697
-
698
- # fill with scalar
699
- fill_value = 42
700
- fa.fill_(fill_value)
701
- assert_np_equal(fa.numpy(), np.full((*fa.shape, *mat_shape), fill_value, dtype=nptype))
702
-
703
- # test zeroing
704
- fa.zero_()
705
- assert_np_equal(fa.numpy(), np.zeros((*fa.shape, *mat_shape), dtype=nptype))
706
-
707
- # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
708
- if wptype != wp.bool:
709
- fill_arr1 = np.arange(mat_len, dtype=nptype)
710
- else:
711
- fill_arr1 = np.ones(mat_len, dtype=nptype)
712
-
713
- fill_arr2 = fill_arr1.reshape(mat_shape)
714
- fill_list1 = list(fill_arr1)
715
- fill_list2 = [list(row) for row in fill_arr2]
716
- fill_mat = mat_type(fill_arr1)
717
-
718
- expected = np.tile(fill_arr1, fa.size).reshape((*fa.shape, *mat_shape))
719
-
720
- # fill with 1d numpy array
721
- fa.fill_(fill_arr1)
722
- assert_np_equal(fa.numpy(), expected)
723
-
724
- # clear
725
- fa.zero_()
726
-
727
- # fill with 2d numpy array
728
- fa.fill_(fill_arr2)
729
- assert_np_equal(fa.numpy(), expected)
730
-
731
- # clear
732
- fa.zero_()
733
-
734
- # fill with flat list
735
- fa.fill_(fill_list1)
736
- assert_np_equal(fa.numpy(), expected)
737
-
738
- # clear
739
- fa.zero_()
740
-
741
- # fill with nested list
742
- fa.fill_(fill_list2)
743
- assert_np_equal(fa.numpy(), expected)
744
-
745
- # clear
746
- fa.zero_()
747
-
748
- # fill with mat instance
749
- fa.fill_(fill_mat)
750
- assert_np_equal(fa.numpy(), expected)
751
-
752
- # reset data
753
- wp.copy(fa, data)
754
-
755
- # test indexed
756
- indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
757
- ifa = fa[indices1]
758
-
759
- # ensure that the other indices remain unchanged
760
- indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
761
- ifb = fa[indices2]
762
-
763
- assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, *mat_shape), dtype=nptype))
764
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
765
-
766
- # fill with scalar
767
- fill_value = 42
768
- ifa.fill_(fill_value)
769
- assert_np_equal(ifa.numpy(), np.full((*ifa.shape, *mat_shape), fill_value, dtype=nptype))
770
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
771
-
772
- # test zeroing
773
- ifa.zero_()
774
- assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, *mat_shape), dtype=nptype))
775
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
776
-
777
- # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
778
- if wptype != wp.bool:
779
- fill_arr1 = np.arange(mat_len, dtype=nptype)
780
- else:
781
- fill_arr1 = np.ones(mat_len, dtype=nptype)
782
- fill_arr2 = fill_arr1.reshape(mat_shape)
783
- fill_list1 = list(fill_arr1)
784
- fill_list2 = [list(row) for row in fill_arr2]
785
- fill_mat = mat_type(fill_arr1)
786
-
787
- expected = np.tile(fill_arr1, ifa.size).reshape((*ifa.shape, *mat_shape))
788
-
789
- # fill with 1d numpy array
790
- ifa.fill_(fill_arr1)
791
- assert_np_equal(ifa.numpy(), expected)
792
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
793
-
794
- # clear
795
- ifa.zero_()
796
-
797
- # fill with 2d numpy array
798
- ifa.fill_(fill_arr2)
799
- assert_np_equal(ifa.numpy(), expected)
800
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
801
-
802
- # clear
803
- ifa.zero_()
804
-
805
- # fill with flat list
806
- ifa.fill_(fill_list1)
807
- assert_np_equal(ifa.numpy(), expected)
808
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
809
-
810
- # clear
811
- ifa.zero_()
812
-
813
- # fill with nested list
814
- ifa.fill_(fill_list2)
815
- assert_np_equal(ifa.numpy(), expected)
816
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
817
-
818
- # clear
819
- ifa.zero_()
820
-
821
- # fill with mat instance
822
- ifa.fill_(fill_mat)
823
- assert_np_equal(ifa.numpy(), expected)
824
- assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
825
-
826
-
827
- @wp.kernel
828
- def fa_generic_sums_kernel(a: wp.fabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
829
- i = wp.tid()
830
-
831
- # get sub-array using wp::view()
832
- row = a[i]
833
-
834
- # get sub-array length
835
- count = row.shape[0]
836
-
837
- # compute sub-array sum
838
- for j in range(count):
839
- sums[i] = sums[i] + row[j]
840
-
841
-
842
- @wp.kernel
843
- def fa_generic_sums_kernel_indexed(a: wp.indexedfabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
844
- i = wp.tid()
845
-
846
- # get sub-array using wp::view()
847
- row = a[i]
848
-
849
- # get sub-array length
850
- count = row.shape[0]
851
-
852
- # compute sub-array sum
853
- for j in range(count):
854
- sums[i] = sums[i] + row[j]
855
-
856
-
857
- def test_fabricarrayarray(test, device):
858
- for T in _fabric_types:
859
- if hasattr(T, "_wp_scalar_type_"):
860
- nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
861
- else:
862
- nptype = wp.types.warp_type_to_np_dtype[T]
863
-
864
- n = 100
865
-
866
- min_length = 1
867
- max_length = 10
868
- arrays = []
869
- expected_sums = []
870
- expected_sums_indexed = []
871
-
872
- # generate data arrays
873
- length = min_length
874
- for i in range(n):
875
- if length > max_length:
876
- length = min_length
877
-
878
- na = np.arange(1, length + 1, dtype=nptype)
879
-
880
- arrays.append(wp.array(data=na, device=device))
881
- expected_sums.append(na.sum())
882
-
883
- # every second index
884
- if i % 2 == 0:
885
- expected_sums_indexed.append(na.sum())
886
-
887
- length += 1
888
-
889
- data_iface = _create_fabric_array_array_interface(arrays, "foo")
890
- fa = wp.fabricarrayarray(data=data_iface, attrib="foo")
891
-
892
- sums = wp.zeros_like(fa)
893
-
894
- test.assertEqual(fa.dtype, sums.dtype)
895
- test.assertEqual(fa.ndim, 2)
896
- test.assertEqual(sums.ndim, 1)
897
- test.assertEqual(fa.shape, sums.shape)
898
- test.assertEqual(fa.size, sums.size)
899
-
900
- wp.launch(fa_generic_sums_kernel, dim=fa.size, inputs=[fa, sums], device=device)
901
-
902
- assert_np_equal(sums.numpy(), np.array(expected_sums, dtype=nptype))
903
-
904
- # test indexed
905
- indices = wp.array(data=np.arange(0, n, 2, dtype=np.int32), device=device)
906
- ifa = fa[indices]
907
-
908
- sums = wp.zeros_like(ifa)
909
-
910
- test.assertEqual(ifa.dtype, sums.dtype)
911
- test.assertEqual(ifa.ndim, 2)
912
- test.assertEqual(sums.ndim, 1)
913
- test.assertEqual(ifa.shape, sums.shape)
914
- test.assertEqual(ifa.size, sums.size)
915
-
916
- wp.launch(fa_generic_sums_kernel_indexed, dim=ifa.size, inputs=[ifa, sums], device=device)
917
-
918
- assert_np_equal(sums.numpy(), np.array(expected_sums_indexed, dtype=nptype))
919
-
920
-
921
- # explicit kernel overloads
922
- for T in _fabric_types:
923
- wp.overload(fa_generic_dtype_kernel, [wp.fabricarray(dtype=T), wp.fabricarray(dtype=T)])
924
- wp.overload(fa_generic_dtype_kernel_indexed, [wp.indexedfabricarray(dtype=T), wp.indexedfabricarray(dtype=T)])
925
-
926
- wp.overload(fa_generic_array_kernel, [wp.fabricarray(dtype=T), wp.fabricarray(dtype=T)])
927
- wp.overload(fa_generic_array_kernel, [wp.indexedfabricarray(dtype=T), wp.indexedfabricarray(dtype=T)])
928
-
929
- wp.overload(fa_generic_sums_kernel, [wp.fabricarrayarray(dtype=T), wp.array(dtype=T)])
930
- wp.overload(fa_generic_sums_kernel_indexed, [wp.indexedfabricarrayarray(dtype=T), wp.array(dtype=T)])
931
-
932
-
933
- devices = get_test_devices()
934
-
935
-
936
- class TestFabricArray(unittest.TestCase):
937
- pass
938
-
939
-
940
- # fabric arrays
941
- add_function_test(TestFabricArray, "test_fabricarray_kernel", test_fabricarray_kernel, devices=devices)
942
- add_function_test(TestFabricArray, "test_fabricarray_empty", test_fabricarray_empty, devices=devices)
943
- add_function_test(TestFabricArray, "test_fabricarray_generic_dtype", test_fabricarray_generic_dtype, devices=devices)
944
- add_function_test(TestFabricArray, "test_fabricarray_generic_array", test_fabricarray_generic_array, devices=devices)
945
- add_function_test(TestFabricArray, "test_fabricarray_fill_scalar", test_fabricarray_fill_scalar, devices=devices)
946
- add_function_test(TestFabricArray, "test_fabricarray_fill_vector", test_fabricarray_fill_vector, devices=devices)
947
- add_function_test(TestFabricArray, "test_fabricarray_fill_matrix", test_fabricarray_fill_matrix, devices=devices)
948
-
949
- # fabric arrays of arrays
950
- add_function_test(TestFabricArray, "test_fabricarrayarray", test_fabricarrayarray, devices=devices)
951
-
952
-
953
- if __name__ == "__main__":
954
- wp.build.clear_kernel_cache()
955
- unittest.main(verbosity=2)
1
+ # Copyright (c) 2023 NVIDIA CORPORATION. All rights reserved.
2
+ # NVIDIA CORPORATION and its licensors retain all intellectual property
3
+ # and proprietary rights in and to this software, related documentation
4
+ # and any modifications thereto. Any use, reproduction, disclosure or
5
+ # distribution of this software and related documentation without an express
6
+ # license agreement from NVIDIA CORPORATION is strictly prohibited.
7
+
8
+ import math
9
+ import unittest
10
+ from typing import Any
11
+
12
+ import numpy as np
13
+
14
+ import warp as wp
15
+ from warp.tests.unittest_utils import *
16
+
17
+ # types to test fabric arrays
18
+ _fabric_types = [
19
+ *wp.types.scalar_types,
20
+ *[wp.types.vector(2, T) for T in wp.types.scalar_types],
21
+ *[wp.types.vector(3, T) for T in wp.types.scalar_types],
22
+ *[wp.types.vector(4, T) for T in wp.types.scalar_types],
23
+ *[wp.types.matrix((2, 2), T) for T in wp.types.scalar_types],
24
+ *[wp.types.matrix((3, 3), T) for T in wp.types.scalar_types],
25
+ *[wp.types.matrix((4, 4), T) for T in wp.types.scalar_types],
26
+ *[wp.types.quaternion(T) for T in wp.types.float_types],
27
+ ]
28
+
29
+
30
+ def _warp_type_to_fabric(dtype, is_array=False):
31
+ scalar_map = {
32
+ wp.bool: "b",
33
+ wp.int8: "i1",
34
+ wp.int16: "i2",
35
+ wp.int32: "i4",
36
+ wp.int64: "i8",
37
+ wp.uint8: "u1",
38
+ wp.uint16: "u2",
39
+ wp.uint32: "u4",
40
+ wp.uint64: "u8",
41
+ wp.float16: "f2",
42
+ wp.float32: "f4",
43
+ wp.float64: "f8",
44
+ }
45
+
46
+ if hasattr(dtype, "_wp_scalar_type_"):
47
+ type_str = scalar_map[dtype._wp_scalar_type_]
48
+ if len(dtype._shape_) == 1:
49
+ role = "vector"
50
+ else:
51
+ role = "matrix"
52
+ else:
53
+ type_str = scalar_map[dtype]
54
+ role = ""
55
+
56
+ if is_array:
57
+ array_depth = 1
58
+ else:
59
+ array_depth = 0
60
+
61
+ return (True, type_str, dtype._length_, array_depth, role)
62
+
63
+
64
+ # returns a fabric array interface constructed from a regular array
65
+ def _create_fabric_array_interface(data: wp.array, attrib: str, bucket_sizes: list = None, copy=False):
66
+ assert isinstance(data, wp.array)
67
+ assert data.ndim == 1
68
+
69
+ assert isinstance(attrib, str)
70
+
71
+ if copy:
72
+ data = wp.clone(data)
73
+
74
+ if bucket_sizes is not None:
75
+ assert hasattr(bucket_sizes, "__len__")
76
+
77
+ # verify total size
78
+ total_size = 0
79
+ for bucket_size in bucket_sizes:
80
+ total_size += bucket_size
81
+
82
+ if total_size != data.size:
83
+ raise RuntimeError("Bucket sizes don't add up to the size of data array")
84
+
85
+ elif data.size > 0:
86
+ rng = np.random.default_rng(123)
87
+
88
+ # generate random bucket sizes
89
+ bucket_min = 1
90
+ bucket_max = math.ceil(0.5 * data.size)
91
+ total_size = data.size
92
+ size_remaining = total_size
93
+
94
+ bucket_sizes = []
95
+ while size_remaining >= bucket_max:
96
+ bucket_size = rng.integers(bucket_min, high=bucket_max, dtype=int)
97
+ bucket_sizes.append(bucket_size)
98
+ size_remaining -= bucket_size
99
+
100
+ if size_remaining > 0:
101
+ bucket_sizes.append(size_remaining)
102
+
103
+ else:
104
+ # empty data array
105
+ bucket_sizes = []
106
+
107
+ dtype_size = wp.types.type_size_in_bytes(data.dtype)
108
+ p = int(data.ptr) if data.ptr else 0
109
+ pointers = []
110
+ counts = []
111
+ for bucket_size in bucket_sizes:
112
+ pointers.append(p)
113
+ counts.append(bucket_size)
114
+ p += bucket_size * dtype_size
115
+
116
+ attrib_info = {}
117
+
118
+ attrib_info["type"] = _warp_type_to_fabric(data.dtype)
119
+ attrib_info["access"] = 2 # ReadWrite
120
+ attrib_info["pointers"] = pointers
121
+ attrib_info["counts"] = counts
122
+
123
+ iface = {}
124
+ iface["version"] = 1
125
+ iface["device"] = str(data.device)
126
+ iface["attribs"] = {attrib: attrib_info}
127
+ iface["_ref"] = data # backref to keep the array alive
128
+
129
+ return iface
130
+
131
+
132
+ # returns a fabric array array interface constructed from a list of regular arrays
133
+ def _create_fabric_array_array_interface(data: list, attrib: str, bucket_sizes: list = None):
134
+ # data should be a list of arrays
135
+ assert isinstance(data, list)
136
+
137
+ num_arrays = len(data)
138
+ assert num_arrays > 0
139
+
140
+ device = data[0].device
141
+ dtype = data[0].dtype
142
+
143
+ assert isinstance(attrib, str)
144
+
145
+ if bucket_sizes is not None:
146
+ assert hasattr(bucket_sizes, "__len__")
147
+
148
+ # verify total size
149
+ total_size = 0
150
+ for bucket_size in bucket_sizes:
151
+ total_size += bucket_size
152
+
153
+ if total_size != num_arrays:
154
+ raise RuntimeError("Bucket sizes don't add up to the number of given arrays")
155
+
156
+ else:
157
+ rng = np.random.default_rng(123)
158
+
159
+ # generate random bucket sizes
160
+ bucket_min = 1
161
+ bucket_max = math.ceil(0.5 * num_arrays)
162
+ total_size = num_arrays
163
+ size_remaining = total_size
164
+
165
+ bucket_sizes = []
166
+ while size_remaining >= bucket_max:
167
+ bucket_size = rng.integers(bucket_min, high=bucket_max, dtype=int)
168
+ bucket_sizes.append(bucket_size)
169
+ size_remaining -= bucket_size
170
+
171
+ if size_remaining > 0:
172
+ bucket_sizes.append(size_remaining)
173
+
174
+ # initialize array of pointers to arrays and their lengths
175
+ _array_pointers = []
176
+ _array_lengths = []
177
+ for i in range(num_arrays):
178
+ _array_pointers.append(data[i].ptr)
179
+ _array_lengths.append(data[i].size)
180
+
181
+ array_pointers = wp.array(_array_pointers, dtype=wp.uint64, device=device)
182
+ pointer_size = wp.types.type_size_in_bytes(array_pointers.dtype)
183
+
184
+ lengths = wp.array(_array_lengths, dtype=wp.uint64, device=device)
185
+ length_size = wp.types.type_size_in_bytes(lengths.dtype)
186
+
187
+ p_pointers = int(array_pointers.ptr)
188
+ p_lengths = int(lengths.ptr)
189
+ pointers = []
190
+ counts = []
191
+ array_lengths = []
192
+ for bucket_size in bucket_sizes:
193
+ pointers.append(p_pointers)
194
+ counts.append(bucket_size)
195
+ array_lengths.append(p_lengths)
196
+ p_pointers += bucket_size * pointer_size
197
+ p_lengths += bucket_size * length_size
198
+
199
+ attrib_info = {}
200
+
201
+ attrib_info["type"] = _warp_type_to_fabric(dtype, is_array=True)
202
+ attrib_info["access"] = 2 # ReadWrite
203
+ attrib_info["pointers"] = pointers
204
+ attrib_info["counts"] = counts
205
+ attrib_info["array_lengths"] = array_lengths
206
+
207
+ iface = {}
208
+ iface["version"] = 1
209
+ iface["device"] = str(device)
210
+ iface["attribs"] = {attrib: attrib_info}
211
+ iface["_ref"] = data # backref to keep the data arrays alive
212
+ iface["_ref_pointers"] = array_pointers # backref to keep the array pointers alive
213
+ iface["_ref_lengths"] = lengths # backref to keep the lengths array alive
214
+
215
+ return iface
216
+
217
+
218
+ @wp.kernel
219
+ def fa_kernel(a: wp.fabricarray(dtype=float), expected: wp.array(dtype=float)):
220
+ i = wp.tid()
221
+
222
+ wp.expect_eq(a[i], expected[i])
223
+
224
+ a[i] = 2.0 * a[i]
225
+
226
+ wp.atomic_add(a, i, 1.0)
227
+
228
+ wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)
229
+
230
+
231
+ @wp.kernel
232
+ def fa_kernel_indexed(a: wp.indexedfabricarray(dtype=float), expected: wp.indexedarray(dtype=float)):
233
+ i = wp.tid()
234
+
235
+ wp.expect_eq(a[i], expected[i])
236
+
237
+ a[i] = 2.0 * a[i]
238
+
239
+ wp.atomic_add(a, i, 1.0)
240
+
241
+ wp.expect_eq(a[i], 2.0 * expected[i] + 1.0)
242
+
243
+
244
+ def test_fabricarray_kernel(test, device):
245
+ data = wp.array(data=np.arange(100, dtype=np.float32), device=device)
246
+ iface = _create_fabric_array_interface(data, "foo", copy=True)
247
+ fa = wp.fabricarray(data=iface, attrib="foo")
248
+
249
+ test.assertEqual(fa.dtype, data.dtype)
250
+ test.assertEqual(fa.ndim, 1)
251
+ test.assertEqual(fa.shape, data.shape)
252
+ test.assertEqual(fa.size, data.size)
253
+
254
+ wp.launch(fa_kernel, dim=fa.size, inputs=[fa, data], device=device)
255
+
256
+ # reset data
257
+ wp.copy(fa, data)
258
+
259
+ # test indexed
260
+ indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
261
+ ifa = fa[indices]
262
+ idata = data[indices]
263
+
264
+ test.assertEqual(ifa.dtype, idata.dtype)
265
+ test.assertEqual(ifa.ndim, 1)
266
+ test.assertEqual(ifa.shape, idata.shape)
267
+ test.assertEqual(ifa.size, idata.size)
268
+
269
+ wp.launch(fa_kernel_indexed, dim=ifa.size, inputs=[ifa, idata], device=device)
270
+
271
+ wp.synchronize_device(device)
272
+
273
+
274
+ @wp.kernel
275
+ def fa_generic_dtype_kernel(a: wp.fabricarray(dtype=Any), b: wp.fabricarray(dtype=Any)):
276
+ i = wp.tid()
277
+ b[i] = a[i] + a[i]
278
+
279
+
280
+ @wp.kernel
281
+ def fa_generic_dtype_kernel_indexed(a: wp.indexedfabricarray(dtype=Any), b: wp.indexedfabricarray(dtype=Any)):
282
+ i = wp.tid()
283
+ b[i] = a[i] + a[i]
284
+
285
+
286
+ def test_fabricarray_generic_dtype(test, device):
287
+ for T in _fabric_types:
288
+ if hasattr(T, "_wp_scalar_type_"):
289
+ nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
290
+ else:
291
+ nptype = wp.types.warp_type_to_np_dtype[T]
292
+
293
+ data = wp.array(data=np.arange(10, dtype=nptype), device=device)
294
+ data_iface = _create_fabric_array_interface(data, "foo", copy=True)
295
+ fa = wp.fabricarray(data=data_iface, attrib="foo")
296
+
297
+ result = wp.zeros_like(data)
298
+ result_iface = _create_fabric_array_interface(result, "foo", copy=True)
299
+ fb = wp.fabricarray(data=result_iface, attrib="foo")
300
+
301
+ test.assertEqual(fa.dtype, fb.dtype)
302
+ test.assertEqual(fa.ndim, fb.ndim)
303
+ test.assertEqual(fa.shape, fb.shape)
304
+ test.assertEqual(fa.size, fb.size)
305
+
306
+ wp.launch(fa_generic_dtype_kernel, dim=fa.size, inputs=[fa, fb], device=device)
307
+
308
+ assert_np_equal(fb.numpy(), 2 * fa.numpy())
309
+
310
+ # reset data
311
+ wp.copy(fa, data)
312
+ wp.copy(fb, result)
313
+
314
+ # test indexed
315
+ indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
316
+ ifa = fa[indices]
317
+ ifb = fb[indices]
318
+
319
+ test.assertEqual(ifa.dtype, ifb.dtype)
320
+ test.assertEqual(ifa.ndim, ifb.ndim)
321
+ test.assertEqual(ifa.shape, ifb.shape)
322
+ test.assertEqual(ifa.size, ifb.size)
323
+
324
+ wp.launch(fa_generic_dtype_kernel_indexed, dim=ifa.size, inputs=[ifa, ifb], device=device)
325
+
326
+ assert_np_equal(ifb.numpy(), 2 * ifa.numpy())
327
+
328
+
329
+ @wp.kernel
330
+ def fa_generic_array_kernel(a: Any, b: Any):
331
+ i = wp.tid()
332
+ b[i] = a[i] + a[i]
333
+
334
+
335
+ def test_fabricarray_generic_array(test, device):
336
+ for T in _fabric_types:
337
+ if hasattr(T, "_wp_scalar_type_"):
338
+ nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
339
+ else:
340
+ nptype = wp.types.warp_type_to_np_dtype[T]
341
+
342
+ data = wp.array(data=np.arange(100, dtype=nptype), device=device)
343
+ data_iface = _create_fabric_array_interface(data, "foo", copy=True)
344
+ fa = wp.fabricarray(data=data_iface, attrib="foo")
345
+
346
+ result = wp.zeros_like(data)
347
+ result_iface = _create_fabric_array_interface(result, "foo", copy=True)
348
+ fb = wp.fabricarray(data=result_iface, attrib="foo")
349
+
350
+ test.assertEqual(fa.dtype, fb.dtype)
351
+ test.assertEqual(fa.ndim, fb.ndim)
352
+ test.assertEqual(fa.shape, fb.shape)
353
+ test.assertEqual(fa.size, fb.size)
354
+
355
+ wp.launch(fa_generic_array_kernel, dim=fa.size, inputs=[fa, fb], device=device)
356
+
357
+ assert_np_equal(fb.numpy(), 2 * fa.numpy())
358
+
359
+ # reset data
360
+ wp.copy(fa, data)
361
+ wp.copy(fb, result)
362
+
363
+ # test indexed
364
+ indices = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
365
+ ifa = fa[indices]
366
+ ifb = fb[indices]
367
+
368
+ test.assertEqual(ifa.dtype, ifb.dtype)
369
+ test.assertEqual(ifa.ndim, ifb.ndim)
370
+ test.assertEqual(ifa.shape, ifb.shape)
371
+ test.assertEqual(ifa.size, ifb.size)
372
+
373
+ wp.launch(fa_generic_array_kernel, dim=ifa.size, inputs=[ifa, ifb], device=device)
374
+
375
+ assert_np_equal(ifb.numpy(), 2 * ifa.numpy())
376
+
377
+
378
+ def test_fabricarray_empty(test, device):
379
+ # Test whether common operations work with empty (zero-sized) indexed arrays
380
+ # without throwing exceptions.
381
+
382
+ def test_empty_ops(nrows, ncols, wptype, nptype):
383
+ # scalar, vector, or matrix
384
+ if ncols > 0:
385
+ if nrows > 0:
386
+ wptype = wp.types.matrix((nrows, ncols), wptype)
387
+ else:
388
+ wptype = wp.types.vector(ncols, wptype)
389
+ dtype_shape = wptype._shape_
390
+ else:
391
+ dtype_shape = ()
392
+
393
+ fill_value = wptype(42)
394
+
395
+ # create an empty data array
396
+ data = wp.empty(0, dtype=wptype, device=device)
397
+ iface = _create_fabric_array_interface(data, "foo", copy=True)
398
+ fa = wp.fabricarray(data=iface, attrib="foo")
399
+
400
+ test.assertEqual(fa.size, 0)
401
+ test.assertEqual(fa.shape, (0,))
402
+
403
+ # all of these methods should succeed with zero-sized arrays
404
+ fa.zero_()
405
+ fa.fill_(fill_value)
406
+ fb = fa.contiguous()
407
+
408
+ fb = wp.empty_like(fa)
409
+ fb = wp.zeros_like(fa)
410
+ fb = wp.full_like(fa, fill_value)
411
+ fb = wp.clone(fa)
412
+
413
+ wp.copy(fa, fb)
414
+ fa.assign(fb)
415
+
416
+ na = fa.numpy()
417
+ test.assertEqual(na.size, 0)
418
+ test.assertEqual(na.shape, (0, *dtype_shape))
419
+ test.assertEqual(na.dtype, nptype)
420
+
421
+ test.assertEqual(fa.list(), [])
422
+
423
+ # test indexed
424
+
425
+ # create a zero-sized array of indices
426
+ indices = wp.empty(0, dtype=int, device=device)
427
+
428
+ ifa = fa[indices]
429
+
430
+ test.assertEqual(ifa.size, 0)
431
+ test.assertEqual(ifa.shape, (0,))
432
+
433
+ # all of these methods should succeed with zero-sized arrays
434
+ ifa.zero_()
435
+ ifa.fill_(fill_value)
436
+ ifb = ifa.contiguous()
437
+
438
+ ifb = wp.empty_like(ifa)
439
+ ifb = wp.zeros_like(ifa)
440
+ ifb = wp.full_like(ifa, fill_value)
441
+ ifb = wp.clone(ifa)
442
+
443
+ wp.copy(ifa, ifb)
444
+ ifa.assign(ifb)
445
+
446
+ na = ifa.numpy()
447
+ test.assertEqual(na.size, 0)
448
+ test.assertEqual(na.shape, (0, *dtype_shape))
449
+ test.assertEqual(na.dtype, nptype)
450
+
451
+ test.assertEqual(ifa.list(), [])
452
+
453
+ # test with scalars, vectors, and matrices
454
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
455
+ # scalars
456
+ test_empty_ops(0, 0, wptype, nptype)
457
+
458
+ for ncols in [2, 3, 4, 5]:
459
+ # vectors
460
+ test_empty_ops(0, ncols, wptype, nptype)
461
+ # square matrices (the Fabric interface only supports square matrices right now)
462
+ test_empty_ops(ncols, ncols, wptype, nptype)
463
+
464
+
465
+ def test_fabricarray_fill_scalar(test, device):
466
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
467
+ # create a data array
468
+ data = wp.zeros(100, dtype=wptype, device=device)
469
+ iface = _create_fabric_array_interface(data, "foo", copy=True)
470
+ fa = wp.fabricarray(data=iface, attrib="foo")
471
+
472
+ assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))
473
+
474
+ # fill with int value
475
+ fill_value = 42
476
+ fa.fill_(fill_value)
477
+ assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value, dtype=nptype))
478
+
479
+ fa.zero_()
480
+ assert_np_equal(fa.numpy(), np.zeros(fa.shape, dtype=nptype))
481
+
482
+ if wptype in wp.types.float_types:
483
+ # fill with float value
484
+ fill_value = 13.37
485
+ fa.fill_(fill_value)
486
+ assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value, dtype=nptype))
487
+
488
+ # fill with Warp scalar value
489
+ fill_value = wptype(17)
490
+ fa.fill_(fill_value)
491
+ assert_np_equal(fa.numpy(), np.full(fa.shape, fill_value.value, dtype=nptype))
492
+
493
+ # reset data
494
+ wp.copy(fa, data)
495
+
496
+ # test indexed
497
+ indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
498
+ ifa = fa[indices1]
499
+
500
+ # ensure that the other indices remain unchanged
501
+ indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
502
+ ifb = fa[indices2]
503
+
504
+ assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
505
+ assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
506
+
507
+ # fill with int value
508
+ fill_value = 42
509
+ ifa.fill_(fill_value)
510
+ assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value, dtype=nptype))
511
+ assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
512
+
513
+ ifa.zero_()
514
+ assert_np_equal(ifa.numpy(), np.zeros(ifa.shape, dtype=nptype))
515
+ assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
516
+
517
+ if wptype in wp.types.float_types:
518
+ # fill with float value
519
+ fill_value = 13.37
520
+ ifa.fill_(fill_value)
521
+ assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value, dtype=nptype))
522
+ assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
523
+
524
+ # fill with Warp scalar value
525
+ fill_value = wptype(17)
526
+ ifa.fill_(fill_value)
527
+ assert_np_equal(ifa.numpy(), np.full(ifa.shape, fill_value.value, dtype=nptype))
528
+ assert_np_equal(ifb.numpy(), np.zeros(ifb.shape, dtype=nptype))
529
+
530
+
531
+ def test_fabricarray_fill_vector(test, device):
532
+ # test filling a vector array with scalar or vector values (vec_type, list, or numpy array)
533
+
534
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
535
+ # vector types
536
+ vector_types = [
537
+ wp.types.vector(2, wptype),
538
+ wp.types.vector(3, wptype),
539
+ wp.types.vector(4, wptype),
540
+ wp.types.vector(5, wptype),
541
+ ]
542
+
543
+ for vec_type in vector_types:
544
+ vec_len = vec_type._length_
545
+
546
+ data = wp.zeros(100, dtype=vec_type, device=device)
547
+ iface = _create_fabric_array_interface(data, "foo", copy=True)
548
+ fa = wp.fabricarray(data=iface, attrib="foo")
549
+
550
+ assert_np_equal(fa.numpy(), np.zeros((*fa.shape, vec_len), dtype=nptype))
551
+
552
+ # fill with int scalar
553
+ fill_value = 42
554
+ fa.fill_(fill_value)
555
+ assert_np_equal(fa.numpy(), np.full((*fa.shape, vec_len), fill_value, dtype=nptype))
556
+
557
+ # test zeroing
558
+ fa.zero_()
559
+ assert_np_equal(fa.numpy(), np.zeros((*fa.shape, vec_len), dtype=nptype))
560
+
561
+ # vector values can be passed as a list, numpy array, or Warp vector instance
562
+ fill_list = [17, 42, 99, 101, 127][:vec_len]
563
+ fill_arr = np.array(fill_list, dtype=nptype)
564
+ fill_vec = vec_type(fill_list)
565
+
566
+ expected = np.tile(fill_arr, fa.size).reshape((*fa.shape, vec_len))
567
+
568
+ # fill with list of vector length
569
+ fa.fill_(fill_list)
570
+ assert_np_equal(fa.numpy(), expected)
571
+
572
+ # clear
573
+ fa.zero_()
574
+
575
+ # fill with numpy array of vector length
576
+ fa.fill_(fill_arr)
577
+ assert_np_equal(fa.numpy(), expected)
578
+
579
+ # clear
580
+ fa.zero_()
581
+
582
+ # fill with vec instance
583
+ fa.fill_(fill_vec)
584
+ assert_np_equal(fa.numpy(), expected)
585
+
586
+ if wptype in wp.types.float_types:
587
+ # fill with float scalar
588
+ fill_value = 13.37
589
+ fa.fill_(fill_value)
590
+ assert_np_equal(fa.numpy(), np.full((*fa.shape, vec_len), fill_value, dtype=nptype))
591
+
592
+ # fill with float list of vector length
593
+ fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
594
+
595
+ fa.fill_(fill_list)
596
+
597
+ expected = np.tile(np.array(fill_list, dtype=nptype), fa.size).reshape((*fa.shape, vec_len))
598
+
599
+ assert_np_equal(fa.numpy(), expected)
600
+
601
+ # reset data
602
+ wp.copy(fa, data)
603
+
604
+ # test indexed
605
+ indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
606
+ ifa = fa[indices1]
607
+
608
+ # ensure that the other indices remain unchanged
609
+ indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
610
+ ifb = fa[indices2]
611
+
612
+ assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, vec_len), dtype=nptype))
613
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
614
+
615
+ # fill with int scalar
616
+ fill_value = 42
617
+ ifa.fill_(fill_value)
618
+ assert_np_equal(ifa.numpy(), np.full((*ifa.shape, vec_len), fill_value, dtype=nptype))
619
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
620
+
621
+ # test zeroing
622
+ ifa.zero_()
623
+ assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, vec_len), dtype=nptype))
624
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
625
+
626
+ # vector values can be passed as a list, numpy array, or Warp vector instance
627
+ fill_list = [17, 42, 99, 101, 127][:vec_len]
628
+ fill_arr = np.array(fill_list, dtype=nptype)
629
+ fill_vec = vec_type(fill_list)
630
+
631
+ expected = np.tile(fill_arr, ifa.size).reshape((*ifa.shape, vec_len))
632
+
633
+ # fill with list of vector length
634
+ ifa.fill_(fill_list)
635
+ assert_np_equal(ifa.numpy(), expected)
636
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
637
+
638
+ # clear
639
+ ifa.zero_()
640
+
641
+ # fill with numpy array of vector length
642
+ ifa.fill_(fill_arr)
643
+ assert_np_equal(ifa.numpy(), expected)
644
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
645
+
646
+ # clear
647
+ ifa.zero_()
648
+
649
+ # fill with vec instance
650
+ ifa.fill_(fill_vec)
651
+ assert_np_equal(ifa.numpy(), expected)
652
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
653
+
654
+ if wptype in wp.types.float_types:
655
+ # fill with float scalar
656
+ fill_value = 13.37
657
+ ifa.fill_(fill_value)
658
+ assert_np_equal(ifa.numpy(), np.full((*ifa.shape, vec_len), fill_value, dtype=nptype))
659
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
660
+
661
+ # fill with float list of vector length
662
+ fill_list = [-2.5, -1.25, 1.25, 2.5, 5.0][:vec_len]
663
+
664
+ ifa.fill_(fill_list)
665
+
666
+ expected = np.tile(np.array(fill_list, dtype=nptype), ifa.size).reshape((*ifa.shape, vec_len))
667
+
668
+ assert_np_equal(ifa.numpy(), expected)
669
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, vec_len), dtype=nptype))
670
+
671
+
672
+ def test_fabricarray_fill_matrix(test, device):
673
+ # test filling a matrix array with scalar or matrix values (mat_type, nested list, or 2d numpy array)
674
+
675
+ for nptype, wptype in wp.types.np_dtype_to_warp_type.items():
676
+ # matrix types
677
+ matrix_types = [
678
+ # square matrices only
679
+ wp.types.matrix((2, 2), wptype),
680
+ wp.types.matrix((3, 3), wptype),
681
+ wp.types.matrix((4, 4), wptype),
682
+ wp.types.matrix((5, 5), wptype),
683
+ ]
684
+
685
+ for mat_type in matrix_types:
686
+ mat_len = mat_type._length_
687
+ mat_shape = mat_type._shape_
688
+
689
+ data = wp.zeros(100, dtype=mat_type, device=device)
690
+ iface = _create_fabric_array_interface(data, "foo", copy=True)
691
+ fa = wp.fabricarray(data=iface, attrib="foo")
692
+
693
+ assert_np_equal(fa.numpy(), np.zeros((*fa.shape, *mat_shape), dtype=nptype))
694
+
695
+ # fill with scalar
696
+ fill_value = 42
697
+ fa.fill_(fill_value)
698
+ assert_np_equal(fa.numpy(), np.full((*fa.shape, *mat_shape), fill_value, dtype=nptype))
699
+
700
+ # test zeroing
701
+ fa.zero_()
702
+ assert_np_equal(fa.numpy(), np.zeros((*fa.shape, *mat_shape), dtype=nptype))
703
+
704
+ # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
705
+ if wptype != wp.bool:
706
+ fill_arr1 = np.arange(mat_len, dtype=nptype)
707
+ else:
708
+ fill_arr1 = np.ones(mat_len, dtype=nptype)
709
+
710
+ fill_arr2 = fill_arr1.reshape(mat_shape)
711
+ fill_list1 = list(fill_arr1)
712
+ fill_list2 = [list(row) for row in fill_arr2]
713
+ fill_mat = mat_type(fill_arr1)
714
+
715
+ expected = np.tile(fill_arr1, fa.size).reshape((*fa.shape, *mat_shape))
716
+
717
+ # fill with 1d numpy array
718
+ fa.fill_(fill_arr1)
719
+ assert_np_equal(fa.numpy(), expected)
720
+
721
+ # clear
722
+ fa.zero_()
723
+
724
+ # fill with 2d numpy array
725
+ fa.fill_(fill_arr2)
726
+ assert_np_equal(fa.numpy(), expected)
727
+
728
+ # clear
729
+ fa.zero_()
730
+
731
+ # fill with flat list
732
+ fa.fill_(fill_list1)
733
+ assert_np_equal(fa.numpy(), expected)
734
+
735
+ # clear
736
+ fa.zero_()
737
+
738
+ # fill with nested list
739
+ fa.fill_(fill_list2)
740
+ assert_np_equal(fa.numpy(), expected)
741
+
742
+ # clear
743
+ fa.zero_()
744
+
745
+ # fill with mat instance
746
+ fa.fill_(fill_mat)
747
+ assert_np_equal(fa.numpy(), expected)
748
+
749
+ # reset data
750
+ wp.copy(fa, data)
751
+
752
+ # test indexed
753
+ indices1 = wp.array(data=np.arange(1, data.size, 2, dtype=np.int32), device=device)
754
+ ifa = fa[indices1]
755
+
756
+ # ensure that the other indices remain unchanged
757
+ indices2 = wp.array(data=np.arange(0, data.size, 2, dtype=np.int32), device=device)
758
+ ifb = fa[indices2]
759
+
760
+ assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, *mat_shape), dtype=nptype))
761
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
762
+
763
+ # fill with scalar
764
+ fill_value = 42
765
+ ifa.fill_(fill_value)
766
+ assert_np_equal(ifa.numpy(), np.full((*ifa.shape, *mat_shape), fill_value, dtype=nptype))
767
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
768
+
769
+ # test zeroing
770
+ ifa.zero_()
771
+ assert_np_equal(ifa.numpy(), np.zeros((*ifa.shape, *mat_shape), dtype=nptype))
772
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
773
+
774
+ # matrix values can be passed as a 1d numpy array, 2d numpy array, flat list, nested list, or Warp matrix instance
775
+ if wptype != wp.bool:
776
+ fill_arr1 = np.arange(mat_len, dtype=nptype)
777
+ else:
778
+ fill_arr1 = np.ones(mat_len, dtype=nptype)
779
+ fill_arr2 = fill_arr1.reshape(mat_shape)
780
+ fill_list1 = list(fill_arr1)
781
+ fill_list2 = [list(row) for row in fill_arr2]
782
+ fill_mat = mat_type(fill_arr1)
783
+
784
+ expected = np.tile(fill_arr1, ifa.size).reshape((*ifa.shape, *mat_shape))
785
+
786
+ # fill with 1d numpy array
787
+ ifa.fill_(fill_arr1)
788
+ assert_np_equal(ifa.numpy(), expected)
789
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
790
+
791
+ # clear
792
+ ifa.zero_()
793
+
794
+ # fill with 2d numpy array
795
+ ifa.fill_(fill_arr2)
796
+ assert_np_equal(ifa.numpy(), expected)
797
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
798
+
799
+ # clear
800
+ ifa.zero_()
801
+
802
+ # fill with flat list
803
+ ifa.fill_(fill_list1)
804
+ assert_np_equal(ifa.numpy(), expected)
805
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
806
+
807
+ # clear
808
+ ifa.zero_()
809
+
810
+ # fill with nested list
811
+ ifa.fill_(fill_list2)
812
+ assert_np_equal(ifa.numpy(), expected)
813
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
814
+
815
+ # clear
816
+ ifa.zero_()
817
+
818
+ # fill with mat instance
819
+ ifa.fill_(fill_mat)
820
+ assert_np_equal(ifa.numpy(), expected)
821
+ assert_np_equal(ifb.numpy(), np.zeros((*ifb.shape, *mat_shape), dtype=nptype))
822
+
823
+
824
+ @wp.kernel
825
+ def fa_generic_sums_kernel(a: wp.fabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
826
+ i = wp.tid()
827
+
828
+ # get sub-array using wp::view()
829
+ row = a[i]
830
+
831
+ # get sub-array length
832
+ count = row.shape[0]
833
+
834
+ # compute sub-array sum
835
+ for j in range(count):
836
+ sums[i] = sums[i] + row[j]
837
+
838
+
839
+ @wp.kernel
840
+ def fa_generic_sums_kernel_indexed(a: wp.indexedfabricarrayarray(dtype=Any), sums: wp.array(dtype=Any)):
841
+ i = wp.tid()
842
+
843
+ # get sub-array using wp::view()
844
+ row = a[i]
845
+
846
+ # get sub-array length
847
+ count = row.shape[0]
848
+
849
+ # compute sub-array sum
850
+ for j in range(count):
851
+ sums[i] = sums[i] + row[j]
852
+
853
+
854
+ def test_fabricarrayarray(test, device):
855
+ for T in _fabric_types:
856
+ if hasattr(T, "_wp_scalar_type_"):
857
+ nptype = wp.types.warp_type_to_np_dtype[T._wp_scalar_type_]
858
+ else:
859
+ nptype = wp.types.warp_type_to_np_dtype[T]
860
+
861
+ n = 100
862
+
863
+ min_length = 1
864
+ max_length = 10
865
+ arrays = []
866
+ expected_sums = []
867
+ expected_sums_indexed = []
868
+
869
+ # generate data arrays
870
+ length = min_length
871
+ for i in range(n):
872
+ if length > max_length:
873
+ length = min_length
874
+
875
+ na = np.arange(1, length + 1, dtype=nptype)
876
+
877
+ arrays.append(wp.array(data=na, device=device))
878
+ expected_sums.append(na.sum())
879
+
880
+ # every second index
881
+ if i % 2 == 0:
882
+ expected_sums_indexed.append(na.sum())
883
+
884
+ length += 1
885
+
886
+ data_iface = _create_fabric_array_array_interface(arrays, "foo")
887
+ fa = wp.fabricarrayarray(data=data_iface, attrib="foo")
888
+
889
+ sums = wp.zeros_like(fa)
890
+
891
+ test.assertEqual(fa.dtype, sums.dtype)
892
+ test.assertEqual(fa.ndim, 2)
893
+ test.assertEqual(sums.ndim, 1)
894
+ test.assertEqual(fa.shape, sums.shape)
895
+ test.assertEqual(fa.size, sums.size)
896
+
897
+ wp.launch(fa_generic_sums_kernel, dim=fa.size, inputs=[fa, sums], device=device)
898
+
899
+ assert_np_equal(sums.numpy(), np.array(expected_sums, dtype=nptype))
900
+
901
+ # test indexed
902
+ indices = wp.array(data=np.arange(0, n, 2, dtype=np.int32), device=device)
903
+ ifa = fa[indices]
904
+
905
+ sums = wp.zeros_like(ifa)
906
+
907
+ test.assertEqual(ifa.dtype, sums.dtype)
908
+ test.assertEqual(ifa.ndim, 2)
909
+ test.assertEqual(sums.ndim, 1)
910
+ test.assertEqual(ifa.shape, sums.shape)
911
+ test.assertEqual(ifa.size, sums.size)
912
+
913
+ wp.launch(fa_generic_sums_kernel_indexed, dim=ifa.size, inputs=[ifa, sums], device=device)
914
+
915
+ assert_np_equal(sums.numpy(), np.array(expected_sums_indexed, dtype=nptype))
916
+
917
+
918
+ # explicit kernel overloads
919
+ for T in _fabric_types:
920
+ wp.overload(fa_generic_dtype_kernel, [wp.fabricarray(dtype=T), wp.fabricarray(dtype=T)])
921
+ wp.overload(fa_generic_dtype_kernel_indexed, [wp.indexedfabricarray(dtype=T), wp.indexedfabricarray(dtype=T)])
922
+
923
+ wp.overload(fa_generic_array_kernel, [wp.fabricarray(dtype=T), wp.fabricarray(dtype=T)])
924
+ wp.overload(fa_generic_array_kernel, [wp.indexedfabricarray(dtype=T), wp.indexedfabricarray(dtype=T)])
925
+
926
+ wp.overload(fa_generic_sums_kernel, [wp.fabricarrayarray(dtype=T), wp.array(dtype=T)])
927
+ wp.overload(fa_generic_sums_kernel_indexed, [wp.indexedfabricarrayarray(dtype=T), wp.array(dtype=T)])
928
+
929
+
930
+ devices = get_test_devices()
931
+
932
+
933
+ class TestFabricArray(unittest.TestCase):
934
+ pass
935
+
936
+
937
+ # fabric arrays
938
+ add_function_test(TestFabricArray, "test_fabricarray_kernel", test_fabricarray_kernel, devices=devices)
939
+ add_function_test(TestFabricArray, "test_fabricarray_empty", test_fabricarray_empty, devices=devices)
940
+ add_function_test(TestFabricArray, "test_fabricarray_generic_dtype", test_fabricarray_generic_dtype, devices=devices)
941
+ add_function_test(TestFabricArray, "test_fabricarray_generic_array", test_fabricarray_generic_array, devices=devices)
942
+ add_function_test(TestFabricArray, "test_fabricarray_fill_scalar", test_fabricarray_fill_scalar, devices=devices)
943
+ add_function_test(TestFabricArray, "test_fabricarray_fill_vector", test_fabricarray_fill_vector, devices=devices)
944
+ add_function_test(TestFabricArray, "test_fabricarray_fill_matrix", test_fabricarray_fill_matrix, devices=devices)
945
+
946
+ # fabric arrays of arrays
947
+ add_function_test(TestFabricArray, "test_fabricarrayarray", test_fabricarrayarray, devices=devices)
948
+
949
+
950
+ if __name__ == "__main__":
951
+ wp.build.clear_kernel_cache()
952
+ unittest.main(verbosity=2)