warp-lang 1.0.2__py3-none-win_amd64.whl → 1.1.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 (346) 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 +115 -113
  6. warp/build_dll.py +383 -375
  7. warp/builtins.py +3425 -3354
  8. warp/codegen.py +2878 -2792
  9. warp/config.py +40 -36
  10. warp/constants.py +45 -45
  11. warp/context.py +5194 -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 +383 -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 +146 -146
  34. warp/examples/benchmarks/benchmark_launches.py +295 -295
  35. warp/examples/browse.py +29 -29
  36. warp/examples/core/example_dem.py +234 -219
  37. warp/examples/core/example_fluid.py +293 -267
  38. warp/examples/core/example_graph_capture.py +144 -126
  39. warp/examples/core/example_marching_cubes.py +188 -174
  40. warp/examples/core/example_mesh.py +174 -155
  41. warp/examples/core/example_mesh_intersect.py +205 -193
  42. warp/examples/core/example_nvdb.py +176 -170
  43. warp/examples/core/example_raycast.py +105 -90
  44. warp/examples/core/example_raymarch.py +199 -178
  45. warp/examples/core/example_render_opengl.py +185 -141
  46. warp/examples/core/example_sph.py +405 -387
  47. warp/examples/core/example_torch.py +222 -181
  48. warp/examples/core/example_wave.py +263 -248
  49. warp/examples/fem/bsr_utils.py +378 -380
  50. warp/examples/fem/example_apic_fluid.py +407 -389
  51. warp/examples/fem/example_convection_diffusion.py +182 -168
  52. warp/examples/fem/example_convection_diffusion_dg.py +219 -209
  53. warp/examples/fem/example_convection_diffusion_dg0.py +204 -194
  54. warp/examples/fem/example_deformed_geometry.py +177 -159
  55. warp/examples/fem/example_diffusion.py +201 -173
  56. warp/examples/fem/example_diffusion_3d.py +177 -152
  57. warp/examples/fem/example_diffusion_mgpu.py +221 -214
  58. warp/examples/fem/example_mixed_elasticity.py +244 -222
  59. warp/examples/fem/example_navier_stokes.py +259 -243
  60. warp/examples/fem/example_stokes.py +220 -192
  61. warp/examples/fem/example_stokes_transfer.py +265 -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 +260 -246
  65. warp/examples/optim/example_cloth_throw.py +222 -209
  66. warp/examples/optim/example_diffray.py +566 -536
  67. warp/examples/optim/example_drone.py +864 -835
  68. warp/examples/optim/example_inverse_kinematics.py +176 -168
  69. warp/examples/optim/example_inverse_kinematics_torch.py +185 -169
  70. warp/examples/optim/example_spring_cage.py +239 -231
  71. warp/examples/optim/example_trajectory.py +223 -199
  72. warp/examples/optim/example_walker.py +306 -293
  73. warp/examples/sim/example_cartpole.py +139 -129
  74. warp/examples/sim/example_cloth.py +196 -186
  75. warp/examples/sim/example_granular.py +124 -111
  76. warp/examples/sim/example_granular_collision_sdf.py +197 -186
  77. warp/examples/sim/example_jacobian_ik.py +236 -214
  78. warp/examples/sim/example_particle_chain.py +118 -105
  79. warp/examples/sim/example_quadruped.py +193 -180
  80. warp/examples/sim/example_rigid_chain.py +197 -187
  81. warp/examples/sim/example_rigid_contact.py +189 -177
  82. warp/examples/sim/example_rigid_force.py +127 -125
  83. warp/examples/sim/example_rigid_gyroscopic.py +109 -95
  84. warp/examples/sim/example_rigid_soft_contact.py +134 -122
  85. warp/examples/sim/example_soft_body.py +190 -177
  86. warp/fabric.py +337 -335
  87. warp/fem/__init__.py +60 -27
  88. warp/fem/cache.py +401 -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 +15 -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 +744 -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 +441 -435
  104. warp/fem/geometry/hexmesh.py +953 -953
  105. warp/fem/geometry/partition.py +374 -376
  106. warp/fem/geometry/quadmesh_2d.py +532 -532
  107. warp/fem/geometry/tetmesh.py +840 -840
  108. warp/fem/geometry/trimesh_2d.py +577 -577
  109. warp/fem/integrate.py +1630 -1615
  110. warp/fem/operator.py +190 -191
  111. warp/fem/polynomial.py +214 -213
  112. warp/fem/quadrature/__init__.py +2 -2
  113. warp/fem/quadrature/pic_quadrature.py +243 -245
  114. warp/fem/quadrature/quadrature.py +295 -294
  115. warp/fem/space/__init__.py +294 -292
  116. warp/fem/space/basis_space.py +488 -489
  117. warp/fem/space/collocated_function_space.py +100 -105
  118. warp/fem/space/dof_mapper.py +236 -236
  119. warp/fem/space/function_space.py +148 -145
  120. warp/fem/space/grid_2d_function_space.py +267 -267
  121. warp/fem/space/grid_3d_function_space.py +305 -306
  122. warp/fem/space/hexmesh_function_space.py +350 -352
  123. warp/fem/space/partition.py +350 -350
  124. warp/fem/space/quadmesh_2d_function_space.py +368 -369
  125. warp/fem/space/restriction.py +158 -160
  126. warp/fem/space/shape/__init__.py +13 -15
  127. warp/fem/space/shape/cube_shape_function.py +738 -738
  128. warp/fem/space/shape/shape_function.py +102 -103
  129. warp/fem/space/shape/square_shape_function.py +611 -611
  130. warp/fem/space/shape/tet_shape_function.py +565 -567
  131. warp/fem/space/shape/triangle_shape_function.py +429 -429
  132. warp/fem/space/tetmesh_function_space.py +294 -292
  133. warp/fem/space/topology.py +297 -295
  134. warp/fem/space/trimesh_2d_function_space.py +223 -221
  135. warp/fem/types.py +77 -77
  136. warp/fem/utils.py +495 -495
  137. warp/jax.py +166 -141
  138. warp/jax_experimental.py +341 -339
  139. warp/native/array.h +1072 -1025
  140. warp/native/builtin.h +1560 -1560
  141. warp/native/bvh.cpp +398 -398
  142. warp/native/bvh.cu +525 -525
  143. warp/native/bvh.h +429 -429
  144. warp/native/clang/clang.cpp +495 -464
  145. warp/native/crt.cpp +31 -31
  146. warp/native/crt.h +334 -334
  147. warp/native/cuda_crt.h +1049 -1049
  148. warp/native/cuda_util.cpp +549 -540
  149. warp/native/cuda_util.h +288 -203
  150. warp/native/cutlass_gemm.cpp +34 -34
  151. warp/native/cutlass_gemm.cu +372 -372
  152. warp/native/error.cpp +66 -66
  153. warp/native/error.h +27 -27
  154. warp/native/fabric.h +228 -228
  155. warp/native/hashgrid.cpp +301 -278
  156. warp/native/hashgrid.cu +78 -77
  157. warp/native/hashgrid.h +227 -227
  158. warp/native/initializer_array.h +32 -32
  159. warp/native/intersect.h +1204 -1204
  160. warp/native/intersect_adj.h +365 -365
  161. warp/native/intersect_tri.h +322 -322
  162. warp/native/marching.cpp +2 -2
  163. warp/native/marching.cu +497 -497
  164. warp/native/marching.h +2 -2
  165. warp/native/mat.h +1498 -1498
  166. warp/native/matnn.h +333 -333
  167. warp/native/mesh.cpp +203 -203
  168. warp/native/mesh.cu +293 -293
  169. warp/native/mesh.h +1887 -1887
  170. warp/native/nanovdb/NanoVDB.h +4782 -4782
  171. warp/native/nanovdb/PNanoVDB.h +2553 -2553
  172. warp/native/nanovdb/PNanoVDBWrite.h +294 -294
  173. warp/native/noise.h +850 -850
  174. warp/native/quat.h +1084 -1084
  175. warp/native/rand.h +299 -299
  176. warp/native/range.h +108 -108
  177. warp/native/reduce.cpp +156 -156
  178. warp/native/reduce.cu +348 -348
  179. warp/native/runlength_encode.cpp +61 -61
  180. warp/native/runlength_encode.cu +46 -46
  181. warp/native/scan.cpp +30 -30
  182. warp/native/scan.cu +36 -36
  183. warp/native/scan.h +7 -7
  184. warp/native/solid_angle.h +442 -442
  185. warp/native/sort.cpp +94 -94
  186. warp/native/sort.cu +97 -97
  187. warp/native/sort.h +14 -14
  188. warp/native/sparse.cpp +337 -337
  189. warp/native/sparse.cu +544 -544
  190. warp/native/spatial.h +630 -630
  191. warp/native/svd.h +562 -562
  192. warp/native/temp_buffer.h +30 -30
  193. warp/native/vec.h +1132 -1132
  194. warp/native/volume.cpp +297 -297
  195. warp/native/volume.cu +32 -32
  196. warp/native/volume.h +538 -538
  197. warp/native/volume_builder.cu +425 -425
  198. warp/native/volume_builder.h +19 -19
  199. warp/native/warp.cpp +1057 -1052
  200. warp/native/warp.cu +2943 -2828
  201. warp/native/warp.h +313 -305
  202. warp/optim/__init__.py +9 -9
  203. warp/optim/adam.py +120 -120
  204. warp/optim/linear.py +1104 -939
  205. warp/optim/sgd.py +104 -92
  206. warp/render/__init__.py +10 -10
  207. warp/render/render_opengl.py +3217 -3204
  208. warp/render/render_usd.py +768 -749
  209. warp/render/utils.py +152 -150
  210. warp/sim/__init__.py +52 -59
  211. warp/sim/articulation.py +685 -685
  212. warp/sim/collide.py +1594 -1590
  213. warp/sim/import_mjcf.py +489 -481
  214. warp/sim/import_snu.py +220 -221
  215. warp/sim/import_urdf.py +536 -516
  216. warp/sim/import_usd.py +887 -881
  217. warp/sim/inertia.py +316 -317
  218. warp/sim/integrator.py +234 -233
  219. warp/sim/integrator_euler.py +1956 -1956
  220. warp/sim/integrator_featherstone.py +1910 -1991
  221. warp/sim/integrator_xpbd.py +3294 -3312
  222. warp/sim/model.py +4473 -4314
  223. warp/sim/particles.py +113 -112
  224. warp/sim/render.py +417 -403
  225. warp/sim/utils.py +413 -410
  226. warp/sparse.py +1227 -1227
  227. warp/stubs.py +2109 -2469
  228. warp/tape.py +1162 -225
  229. warp/tests/__init__.py +1 -1
  230. warp/tests/__main__.py +4 -4
  231. warp/tests/assets/torus.usda +105 -105
  232. warp/tests/aux_test_class_kernel.py +26 -26
  233. warp/tests/aux_test_compile_consts_dummy.py +10 -10
  234. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -21
  235. warp/tests/aux_test_dependent.py +22 -22
  236. warp/tests/aux_test_grad_customs.py +23 -23
  237. warp/tests/aux_test_reference.py +11 -11
  238. warp/tests/aux_test_reference_reference.py +10 -10
  239. warp/tests/aux_test_square.py +17 -17
  240. warp/tests/aux_test_unresolved_func.py +14 -14
  241. warp/tests/aux_test_unresolved_symbol.py +14 -14
  242. warp/tests/disabled_kinematics.py +239 -239
  243. warp/tests/run_coverage_serial.py +31 -31
  244. warp/tests/test_adam.py +157 -157
  245. warp/tests/test_arithmetic.py +1124 -1124
  246. warp/tests/test_array.py +2417 -2326
  247. warp/tests/test_array_reduce.py +150 -150
  248. warp/tests/test_async.py +668 -656
  249. warp/tests/test_atomic.py +141 -141
  250. warp/tests/test_bool.py +204 -149
  251. warp/tests/test_builtins_resolution.py +1292 -1292
  252. warp/tests/test_bvh.py +164 -171
  253. warp/tests/test_closest_point_edge_edge.py +228 -228
  254. warp/tests/test_codegen.py +566 -553
  255. warp/tests/test_compile_consts.py +97 -101
  256. warp/tests/test_conditional.py +246 -246
  257. warp/tests/test_copy.py +232 -215
  258. warp/tests/test_ctypes.py +632 -632
  259. warp/tests/test_dense.py +67 -67
  260. warp/tests/test_devices.py +91 -98
  261. warp/tests/test_dlpack.py +530 -529
  262. warp/tests/test_examples.py +400 -378
  263. warp/tests/test_fabricarray.py +955 -955
  264. warp/tests/test_fast_math.py +62 -54
  265. warp/tests/test_fem.py +1277 -1278
  266. warp/tests/test_fp16.py +130 -130
  267. warp/tests/test_func.py +338 -337
  268. warp/tests/test_generics.py +571 -571
  269. warp/tests/test_grad.py +746 -640
  270. warp/tests/test_grad_customs.py +333 -336
  271. warp/tests/test_hash_grid.py +210 -164
  272. warp/tests/test_import.py +39 -39
  273. warp/tests/test_indexedarray.py +1134 -1134
  274. warp/tests/test_intersect.py +67 -67
  275. warp/tests/test_jax.py +307 -307
  276. warp/tests/test_large.py +167 -164
  277. warp/tests/test_launch.py +354 -354
  278. warp/tests/test_lerp.py +261 -261
  279. warp/tests/test_linear_solvers.py +191 -171
  280. warp/tests/test_lvalue.py +421 -493
  281. warp/tests/test_marching_cubes.py +65 -65
  282. warp/tests/test_mat.py +1801 -1827
  283. warp/tests/test_mat_lite.py +115 -115
  284. warp/tests/test_mat_scalar_ops.py +2907 -2889
  285. warp/tests/test_math.py +126 -193
  286. warp/tests/test_matmul.py +500 -499
  287. warp/tests/test_matmul_lite.py +410 -410
  288. warp/tests/test_mempool.py +188 -190
  289. warp/tests/test_mesh.py +284 -324
  290. warp/tests/test_mesh_query_aabb.py +228 -241
  291. warp/tests/test_mesh_query_point.py +692 -702
  292. warp/tests/test_mesh_query_ray.py +292 -303
  293. warp/tests/test_mlp.py +276 -276
  294. warp/tests/test_model.py +110 -110
  295. warp/tests/test_modules_lite.py +39 -39
  296. warp/tests/test_multigpu.py +163 -163
  297. warp/tests/test_noise.py +248 -248
  298. warp/tests/test_operators.py +250 -250
  299. warp/tests/test_options.py +123 -125
  300. warp/tests/test_peer.py +133 -137
  301. warp/tests/test_pinned.py +78 -78
  302. warp/tests/test_print.py +54 -54
  303. warp/tests/test_quat.py +2086 -2086
  304. warp/tests/test_rand.py +288 -288
  305. warp/tests/test_reload.py +217 -217
  306. warp/tests/test_rounding.py +179 -179
  307. warp/tests/test_runlength_encode.py +190 -190
  308. warp/tests/test_sim_grad.py +243 -0
  309. warp/tests/test_sim_kinematics.py +91 -97
  310. warp/tests/test_smoothstep.py +168 -168
  311. warp/tests/test_snippet.py +305 -266
  312. warp/tests/test_sparse.py +468 -460
  313. warp/tests/test_spatial.py +2148 -2148
  314. warp/tests/test_streams.py +486 -473
  315. warp/tests/test_struct.py +710 -675
  316. warp/tests/test_tape.py +173 -148
  317. warp/tests/test_torch.py +743 -743
  318. warp/tests/test_transient_module.py +87 -87
  319. warp/tests/test_types.py +556 -659
  320. warp/tests/test_utils.py +490 -499
  321. warp/tests/test_vec.py +1264 -1268
  322. warp/tests/test_vec_lite.py +73 -73
  323. warp/tests/test_vec_scalar_ops.py +2099 -2099
  324. warp/tests/test_verify_fp.py +94 -94
  325. warp/tests/test_volume.py +737 -736
  326. warp/tests/test_volume_write.py +255 -265
  327. warp/tests/unittest_serial.py +37 -37
  328. warp/tests/unittest_suites.py +363 -359
  329. warp/tests/unittest_utils.py +603 -578
  330. warp/tests/unused_test_misc.py +71 -71
  331. warp/tests/walkthrough_debug.py +85 -85
  332. warp/thirdparty/appdirs.py +598 -598
  333. warp/thirdparty/dlpack.py +143 -143
  334. warp/thirdparty/unittest_parallel.py +566 -561
  335. warp/torch.py +321 -295
  336. warp/types.py +4504 -4450
  337. warp/utils.py +1008 -821
  338. {warp_lang-1.0.2.dist-info → warp_lang-1.1.0.dist-info}/LICENSE.md +126 -126
  339. {warp_lang-1.0.2.dist-info → warp_lang-1.1.0.dist-info}/METADATA +338 -400
  340. warp_lang-1.1.0.dist-info/RECORD +352 -0
  341. warp/examples/assets/cube.usda +0 -42
  342. warp/examples/assets/sphere.usda +0 -56
  343. warp/examples/assets/torus.usda +0 -105
  344. warp_lang-1.0.2.dist-info/RECORD +0 -352
  345. {warp_lang-1.0.2.dist-info → warp_lang-1.1.0.dist-info}/WHEEL +0 -0
  346. {warp_lang-1.0.2.dist-info → warp_lang-1.1.0.dist-info}/top_level.txt +0 -0
@@ -1,955 +1,955 @@
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
+ 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)