warp-lang 0.9.0__py3-none-win_amd64.whl → 0.11.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 (315) hide show
  1. warp/__init__.py +15 -7
  2. warp/__init__.pyi +1 -0
  3. warp/bin/warp-clang.dll +0 -0
  4. warp/bin/warp.dll +0 -0
  5. warp/build.py +22 -443
  6. warp/build_dll.py +384 -0
  7. warp/builtins.py +998 -488
  8. warp/codegen.py +1307 -739
  9. warp/config.py +5 -3
  10. warp/constants.py +6 -0
  11. warp/context.py +1291 -548
  12. warp/dlpack.py +31 -31
  13. warp/fabric.py +326 -0
  14. warp/fem/__init__.py +27 -0
  15. warp/fem/cache.py +389 -0
  16. warp/fem/dirichlet.py +181 -0
  17. warp/fem/domain.py +263 -0
  18. warp/fem/field/__init__.py +101 -0
  19. warp/fem/field/field.py +149 -0
  20. warp/fem/field/nodal_field.py +299 -0
  21. warp/fem/field/restriction.py +21 -0
  22. warp/fem/field/test.py +181 -0
  23. warp/fem/field/trial.py +183 -0
  24. warp/fem/geometry/__init__.py +19 -0
  25. warp/fem/geometry/closest_point.py +70 -0
  26. warp/fem/geometry/deformed_geometry.py +271 -0
  27. warp/fem/geometry/element.py +744 -0
  28. warp/fem/geometry/geometry.py +186 -0
  29. warp/fem/geometry/grid_2d.py +373 -0
  30. warp/fem/geometry/grid_3d.py +435 -0
  31. warp/fem/geometry/hexmesh.py +953 -0
  32. warp/fem/geometry/partition.py +376 -0
  33. warp/fem/geometry/quadmesh_2d.py +532 -0
  34. warp/fem/geometry/tetmesh.py +840 -0
  35. warp/fem/geometry/trimesh_2d.py +577 -0
  36. warp/fem/integrate.py +1616 -0
  37. warp/fem/operator.py +191 -0
  38. warp/fem/polynomial.py +213 -0
  39. warp/fem/quadrature/__init__.py +2 -0
  40. warp/fem/quadrature/pic_quadrature.py +245 -0
  41. warp/fem/quadrature/quadrature.py +294 -0
  42. warp/fem/space/__init__.py +292 -0
  43. warp/fem/space/basis_space.py +489 -0
  44. warp/fem/space/collocated_function_space.py +105 -0
  45. warp/fem/space/dof_mapper.py +236 -0
  46. warp/fem/space/function_space.py +145 -0
  47. warp/fem/space/grid_2d_function_space.py +267 -0
  48. warp/fem/space/grid_3d_function_space.py +306 -0
  49. warp/fem/space/hexmesh_function_space.py +352 -0
  50. warp/fem/space/partition.py +350 -0
  51. warp/fem/space/quadmesh_2d_function_space.py +369 -0
  52. warp/fem/space/restriction.py +160 -0
  53. warp/fem/space/shape/__init__.py +15 -0
  54. warp/fem/space/shape/cube_shape_function.py +738 -0
  55. warp/fem/space/shape/shape_function.py +103 -0
  56. warp/fem/space/shape/square_shape_function.py +611 -0
  57. warp/fem/space/shape/tet_shape_function.py +567 -0
  58. warp/fem/space/shape/triangle_shape_function.py +429 -0
  59. warp/fem/space/tetmesh_function_space.py +292 -0
  60. warp/fem/space/topology.py +295 -0
  61. warp/fem/space/trimesh_2d_function_space.py +221 -0
  62. warp/fem/types.py +77 -0
  63. warp/fem/utils.py +495 -0
  64. warp/native/array.h +164 -55
  65. warp/native/builtin.h +150 -174
  66. warp/native/bvh.cpp +75 -328
  67. warp/native/bvh.cu +406 -23
  68. warp/native/bvh.h +37 -45
  69. warp/native/clang/clang.cpp +136 -24
  70. warp/native/crt.cpp +1 -76
  71. warp/native/crt.h +111 -104
  72. warp/native/cuda_crt.h +1049 -0
  73. warp/native/cuda_util.cpp +15 -3
  74. warp/native/cuda_util.h +3 -1
  75. warp/native/cutlass/tools/library/scripts/conv2d_operation.py +463 -0
  76. warp/native/cutlass/tools/library/scripts/conv3d_operation.py +321 -0
  77. warp/native/cutlass/tools/library/scripts/gemm_operation.py +988 -0
  78. warp/native/cutlass/tools/library/scripts/generator.py +4625 -0
  79. warp/native/cutlass/tools/library/scripts/library.py +799 -0
  80. warp/native/cutlass/tools/library/scripts/manifest.py +402 -0
  81. warp/native/cutlass/tools/library/scripts/pycutlass/docs/source/conf.py +96 -0
  82. warp/native/cutlass/tools/library/scripts/pycutlass/profile/conv/conv2d_f16_sm80.py +106 -0
  83. warp/native/cutlass/tools/library/scripts/pycutlass/profile/gemm/gemm_f32_sm80.py +91 -0
  84. warp/native/cutlass/tools/library/scripts/pycutlass/setup.py +80 -0
  85. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/__init__.py +48 -0
  86. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/arguments.py +118 -0
  87. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/c_types.py +241 -0
  88. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/compiler.py +432 -0
  89. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/conv2d_operation.py +631 -0
  90. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/epilogue.py +1026 -0
  91. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/frontend.py +104 -0
  92. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/gemm_operation.py +1276 -0
  93. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/library.py +744 -0
  94. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/memory_manager.py +74 -0
  95. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/operation.py +110 -0
  96. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/parser.py +619 -0
  97. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/reduction_operation.py +398 -0
  98. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/tensor_ref.py +70 -0
  99. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/__init__.py +4 -0
  100. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/conv2d_testbed.py +646 -0
  101. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/gemm_grouped_testbed.py +235 -0
  102. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/gemm_testbed.py +557 -0
  103. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/test/profiler.py +70 -0
  104. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/type_hint.py +39 -0
  105. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/__init__.py +1 -0
  106. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/device.py +76 -0
  107. warp/native/cutlass/tools/library/scripts/pycutlass/src/pycutlass/utils/reference_model.py +255 -0
  108. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/__init__.py +0 -0
  109. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +201 -0
  110. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +177 -0
  111. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +98 -0
  112. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_dgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +95 -0
  113. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_few_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py +163 -0
  114. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_fixed_channels_f16nhwc_f16nhwc_f16nhwc_tensor_op_f32_sm80.py +187 -0
  115. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +309 -0
  116. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +54 -0
  117. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +96 -0
  118. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_fprop_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +107 -0
  119. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_strided_dgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +253 -0
  120. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f16nhwc_tensor_op_f16_sm80.py +97 -0
  121. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f16nhwc_f16nhwc_f32nhwc_tensor_op_f32_sm80.py +242 -0
  122. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_f32nhwc_f32nhwc_f32nhwc_simt_f32_sm80.py +96 -0
  123. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/conv2d_wgrad_implicit_gemm_tf32nhwc_tf32nhwc_f32nhwc_tensor_op_f32_sm80.py +107 -0
  124. warp/native/cutlass/tools/library/scripts/pycutlass/test/conv/run_all_tests.py +10 -0
  125. warp/native/cutlass/tools/library/scripts/pycutlass/test/frontend/test_frontend.py +146 -0
  126. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/__init__.py +0 -0
  127. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_bf16_sm80.py +96 -0
  128. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f16_sm80.py +447 -0
  129. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f32_sm80.py +146 -0
  130. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_f64_sm80.py +102 -0
  131. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_grouped_sm80.py +203 -0
  132. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/gemm_s8_sm80.py +229 -0
  133. warp/native/cutlass/tools/library/scripts/pycutlass/test/gemm/run_all_tests.py +9 -0
  134. warp/native/cutlass/tools/library/scripts/pycutlass/test/unit/test_sm80.py +453 -0
  135. warp/native/cutlass/tools/library/scripts/rank_2k_operation.py +398 -0
  136. warp/native/cutlass/tools/library/scripts/rank_k_operation.py +387 -0
  137. warp/native/cutlass/tools/library/scripts/rt.py +796 -0
  138. warp/native/cutlass/tools/library/scripts/symm_operation.py +400 -0
  139. warp/native/cutlass/tools/library/scripts/trmm_operation.py +407 -0
  140. warp/native/cutlass_gemm.cu +5 -3
  141. warp/native/exports.h +1240 -949
  142. warp/native/fabric.h +228 -0
  143. warp/native/hashgrid.cpp +4 -4
  144. warp/native/hashgrid.h +22 -2
  145. warp/native/initializer_array.h +2 -2
  146. warp/native/intersect.h +22 -7
  147. warp/native/intersect_adj.h +8 -8
  148. warp/native/intersect_tri.h +13 -16
  149. warp/native/marching.cu +157 -161
  150. warp/native/mat.h +119 -19
  151. warp/native/matnn.h +2 -2
  152. warp/native/mesh.cpp +108 -83
  153. warp/native/mesh.cu +243 -6
  154. warp/native/mesh.h +1547 -458
  155. warp/native/nanovdb/NanoVDB.h +1 -1
  156. warp/native/noise.h +272 -329
  157. warp/native/quat.h +51 -8
  158. warp/native/rand.h +45 -35
  159. warp/native/range.h +6 -2
  160. warp/native/reduce.cpp +157 -0
  161. warp/native/reduce.cu +348 -0
  162. warp/native/runlength_encode.cpp +62 -0
  163. warp/native/runlength_encode.cu +46 -0
  164. warp/native/scan.cu +11 -13
  165. warp/native/scan.h +1 -0
  166. warp/native/solid_angle.h +442 -0
  167. warp/native/sort.cpp +13 -0
  168. warp/native/sort.cu +9 -1
  169. warp/native/sparse.cpp +338 -0
  170. warp/native/sparse.cu +545 -0
  171. warp/native/spatial.h +2 -2
  172. warp/native/temp_buffer.h +30 -0
  173. warp/native/vec.h +126 -24
  174. warp/native/volume.h +120 -0
  175. warp/native/warp.cpp +658 -53
  176. warp/native/warp.cu +660 -68
  177. warp/native/warp.h +112 -12
  178. warp/optim/__init__.py +1 -0
  179. warp/optim/linear.py +922 -0
  180. warp/optim/sgd.py +92 -0
  181. warp/render/render_opengl.py +392 -152
  182. warp/render/render_usd.py +11 -11
  183. warp/sim/__init__.py +2 -2
  184. warp/sim/articulation.py +385 -185
  185. warp/sim/collide.py +21 -8
  186. warp/sim/import_mjcf.py +297 -106
  187. warp/sim/import_urdf.py +389 -210
  188. warp/sim/import_usd.py +198 -97
  189. warp/sim/inertia.py +17 -18
  190. warp/sim/integrator_euler.py +14 -8
  191. warp/sim/integrator_xpbd.py +161 -19
  192. warp/sim/model.py +795 -291
  193. warp/sim/optimizer.py +2 -6
  194. warp/sim/render.py +65 -3
  195. warp/sim/utils.py +3 -0
  196. warp/sparse.py +1227 -0
  197. warp/stubs.py +665 -223
  198. warp/tape.py +66 -15
  199. warp/tests/__main__.py +3 -6
  200. warp/tests/assets/curlnoise_golden.npy +0 -0
  201. warp/tests/assets/pnoise_golden.npy +0 -0
  202. warp/tests/assets/torus.usda +105 -105
  203. warp/tests/{test_class_kernel.py → aux_test_class_kernel.py} +9 -1
  204. warp/tests/aux_test_conditional_unequal_types_kernels.py +21 -0
  205. warp/tests/{test_dependent.py → aux_test_dependent.py} +2 -2
  206. warp/tests/{test_reference.py → aux_test_reference.py} +1 -1
  207. warp/tests/aux_test_unresolved_func.py +14 -0
  208. warp/tests/aux_test_unresolved_symbol.py +14 -0
  209. warp/tests/disabled_kinematics.py +239 -0
  210. warp/tests/run_coverage_serial.py +31 -0
  211. warp/tests/test_adam.py +103 -106
  212. warp/tests/test_arithmetic.py +128 -74
  213. warp/tests/test_array.py +1497 -211
  214. warp/tests/test_array_reduce.py +150 -0
  215. warp/tests/test_atomic.py +64 -28
  216. warp/tests/test_bool.py +99 -0
  217. warp/tests/test_builtins_resolution.py +1292 -0
  218. warp/tests/test_bvh.py +75 -43
  219. warp/tests/test_closest_point_edge_edge.py +54 -57
  220. warp/tests/test_codegen.py +233 -128
  221. warp/tests/test_compile_consts.py +28 -20
  222. warp/tests/test_conditional.py +108 -24
  223. warp/tests/test_copy.py +10 -12
  224. warp/tests/test_ctypes.py +112 -88
  225. warp/tests/test_dense.py +21 -14
  226. warp/tests/test_devices.py +98 -0
  227. warp/tests/test_dlpack.py +136 -108
  228. warp/tests/test_examples.py +277 -0
  229. warp/tests/test_fabricarray.py +955 -0
  230. warp/tests/test_fast_math.py +15 -11
  231. warp/tests/test_fem.py +1271 -0
  232. warp/tests/test_fp16.py +53 -19
  233. warp/tests/test_func.py +187 -74
  234. warp/tests/test_generics.py +194 -49
  235. warp/tests/test_grad.py +180 -116
  236. warp/tests/test_grad_customs.py +176 -0
  237. warp/tests/test_hash_grid.py +52 -37
  238. warp/tests/test_import.py +10 -23
  239. warp/tests/test_indexedarray.py +577 -24
  240. warp/tests/test_intersect.py +18 -9
  241. warp/tests/test_large.py +141 -0
  242. warp/tests/test_launch.py +251 -15
  243. warp/tests/test_lerp.py +64 -65
  244. warp/tests/test_linear_solvers.py +154 -0
  245. warp/tests/test_lvalue.py +493 -0
  246. warp/tests/test_marching_cubes.py +12 -13
  247. warp/tests/test_mat.py +508 -2778
  248. warp/tests/test_mat_lite.py +115 -0
  249. warp/tests/test_mat_scalar_ops.py +2889 -0
  250. warp/tests/test_math.py +103 -9
  251. warp/tests/test_matmul.py +305 -69
  252. warp/tests/test_matmul_lite.py +410 -0
  253. warp/tests/test_mesh.py +71 -14
  254. warp/tests/test_mesh_query_aabb.py +41 -25
  255. warp/tests/test_mesh_query_point.py +325 -34
  256. warp/tests/test_mesh_query_ray.py +39 -22
  257. warp/tests/test_mlp.py +30 -22
  258. warp/tests/test_model.py +92 -89
  259. warp/tests/test_modules_lite.py +39 -0
  260. warp/tests/test_multigpu.py +88 -114
  261. warp/tests/test_noise.py +12 -11
  262. warp/tests/test_operators.py +16 -20
  263. warp/tests/test_options.py +11 -11
  264. warp/tests/test_pinned.py +17 -18
  265. warp/tests/test_print.py +32 -11
  266. warp/tests/test_quat.py +275 -129
  267. warp/tests/test_rand.py +18 -16
  268. warp/tests/test_reload.py +38 -34
  269. warp/tests/test_rounding.py +50 -43
  270. warp/tests/test_runlength_encode.py +190 -0
  271. warp/tests/test_smoothstep.py +9 -11
  272. warp/tests/test_snippet.py +143 -0
  273. warp/tests/test_sparse.py +460 -0
  274. warp/tests/test_spatial.py +276 -243
  275. warp/tests/test_streams.py +110 -85
  276. warp/tests/test_struct.py +331 -85
  277. warp/tests/test_tape.py +39 -21
  278. warp/tests/test_torch.py +118 -89
  279. warp/tests/test_transient_module.py +12 -13
  280. warp/tests/test_types.py +614 -0
  281. warp/tests/test_utils.py +494 -0
  282. warp/tests/test_vec.py +354 -1987
  283. warp/tests/test_vec_lite.py +73 -0
  284. warp/tests/test_vec_scalar_ops.py +2099 -0
  285. warp/tests/test_volume.py +457 -293
  286. warp/tests/test_volume_write.py +124 -134
  287. warp/tests/unittest_serial.py +35 -0
  288. warp/tests/unittest_suites.py +341 -0
  289. warp/tests/unittest_utils.py +568 -0
  290. warp/tests/unused_test_misc.py +71 -0
  291. warp/tests/{test_debug.py → walkthough_debug.py} +3 -17
  292. warp/thirdparty/appdirs.py +36 -45
  293. warp/thirdparty/unittest_parallel.py +549 -0
  294. warp/torch.py +72 -30
  295. warp/types.py +1744 -713
  296. warp/utils.py +360 -350
  297. warp_lang-0.11.0.dist-info/LICENSE.md +36 -0
  298. warp_lang-0.11.0.dist-info/METADATA +238 -0
  299. warp_lang-0.11.0.dist-info/RECORD +332 -0
  300. {warp_lang-0.9.0.dist-info → warp_lang-0.11.0.dist-info}/WHEEL +1 -1
  301. warp/bin/warp-clang.exp +0 -0
  302. warp/bin/warp-clang.lib +0 -0
  303. warp/bin/warp.exp +0 -0
  304. warp/bin/warp.lib +0 -0
  305. warp/tests/test_all.py +0 -215
  306. warp/tests/test_array_scan.py +0 -60
  307. warp/tests/test_base.py +0 -208
  308. warp/tests/test_unresolved_func.py +0 -7
  309. warp/tests/test_unresolved_symbol.py +0 -7
  310. warp_lang-0.9.0.dist-info/METADATA +0 -20
  311. warp_lang-0.9.0.dist-info/RECORD +0 -177
  312. /warp/tests/{test_compile_consts_dummy.py → aux_test_compile_consts_dummy.py} +0 -0
  313. /warp/tests/{test_reference_reference.py → aux_test_reference_reference.py} +0 -0
  314. /warp/tests/{test_square.py → aux_test_square.py} +0 -0
  315. {warp_lang-0.9.0.dist-info → warp_lang-0.11.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,447 @@
1
+ import pycutlass
2
+ from pycutlass import *
3
+ from pycutlass.test import *
4
+ import unittest
5
+
6
+ from pycutlass.test.gemm_testbed import test_all_gemm
7
+ from pycutlass.utils.device import device_cc
8
+
9
+
10
+ @unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
11
+ class GemmF16Sm80(unittest.TestCase):
12
+ def test_SM80_Device_Gemm_f32t_f32n_f32t_tensor_op_bf16_f32_128x128x32_64x64x32(self):
13
+ math_inst = MathInstruction(
14
+ instruction_shape=[16, 8, 16],
15
+ element_a=cutlass.float16, element_b=cutlass.float16,
16
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
17
+ math_operation=MathOperation.multiply_add
18
+ )
19
+
20
+ tile_description = TileDescription(
21
+ threadblock_shape=[128, 128, 32],
22
+ stages=3, warp_count=[2, 2, 1],
23
+ math_instruction=math_inst
24
+ )
25
+
26
+ A = TensorDescription(
27
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
28
+ alignment=8
29
+ )
30
+ B = TensorDescription(
31
+ element=cutlass.float16, layout=cutlass.RowMajor,
32
+ alignment=8
33
+ )
34
+ C = TensorDescription(
35
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
36
+ alignment=4
37
+ )
38
+
39
+ element_epilogue = cutlass.float32
40
+
41
+ epilogue_functor = LinearCombination(
42
+ C.element, C.alignment,
43
+ math_inst.element_accumulator, element_epilogue)
44
+
45
+ swizzling_functor = cutlass.BatchedIdentitySwizzle
46
+
47
+ operation = GemmOperationUniversal(
48
+ arch=80, tile_description=tile_description,
49
+ A=A, B=B, C=C,
50
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor,
51
+ direct_store=True
52
+ )
53
+
54
+ self.assertTrue(test_all_gemm(operation, "universal"))
55
+
56
+ def test_SM80_Device_Gemm_f16n_f16n_f16t_tensor_op_f32_128x128x64_64x64x64(self):
57
+ math_inst = MathInstruction(
58
+ instruction_shape=[16, 8, 16],
59
+ element_a=cutlass.float16, element_b=cutlass.float16,
60
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
61
+ math_operation=MathOperation.multiply_add
62
+ )
63
+
64
+ tile_description = TileDescription(
65
+ threadblock_shape=[128, 128, 64],
66
+ stages=3, warp_count=[2, 2, 1],
67
+ math_instruction=math_inst
68
+ )
69
+
70
+ A = TensorDescription(
71
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
72
+ alignment=8
73
+ )
74
+ B = TensorDescription(
75
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
76
+ alignment=8
77
+ )
78
+ C = TensorDescription(
79
+ element=cutlass.float16, layout=cutlass.RowMajor,
80
+ alignment=8
81
+ )
82
+
83
+ element_epilogue = cutlass.float32
84
+
85
+ epilogue_functor = LinearCombination(
86
+ C.element, C.alignment,
87
+ math_inst.element_accumulator, element_epilogue)
88
+
89
+ swizzling_functor = cutlass.IdentitySwizzle1
90
+
91
+ operation = GemmOperationUniversal(
92
+ arch=80, tile_description=tile_description,
93
+ A=A, B=B, C=C,
94
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
95
+ )
96
+
97
+ self.assertTrue(test_all_gemm(operation, "universal"))
98
+
99
+ def test_SM80_Device_Gemm_f16n_f16n_f32n_tensor_op_f32_128x256x64_64x64x64(self):
100
+ math_inst = MathInstruction(
101
+ instruction_shape=[16, 8, 16],
102
+ element_a=cutlass.float16, element_b=cutlass.float16,
103
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
104
+ math_operation=MathOperation.multiply_add
105
+ )
106
+
107
+ tile_description = TileDescription(
108
+ threadblock_shape=[128, 256, 64],
109
+ stages=3, warp_count=[2, 4, 1],
110
+ math_instruction=math_inst
111
+ )
112
+
113
+ A = TensorDescription(
114
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
115
+ alignment=8
116
+ )
117
+ B = TensorDescription(
118
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
119
+ alignment=8
120
+ )
121
+ C = TensorDescription(
122
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
123
+ alignment=4
124
+ )
125
+
126
+ element_epilogue = cutlass.float32
127
+
128
+ epilogue_functor = LinearCombination(
129
+ C.element, C.alignment,
130
+ math_inst.element_accumulator, element_epilogue)
131
+
132
+ swizzling_functor = cutlass.IdentitySwizzle1
133
+
134
+ operation = GemmOperationUniversal(
135
+ arch=80, tile_description=tile_description,
136
+ A=A, B=B, C=C,
137
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
138
+ )
139
+
140
+ self.assertTrue(test_all_gemm(operation, "universal"))
141
+
142
+ def test_SM80_Device_Gemm_f16n_f16n_f32t_tensor_op_f32_256x128x64_64x64x64(self):
143
+ math_inst = MathInstruction(
144
+ instruction_shape=[16, 8, 16],
145
+ element_a=cutlass.float16, element_b=cutlass.float16,
146
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
147
+ math_operation=MathOperation.multiply_add
148
+ )
149
+
150
+ tile_description = TileDescription(
151
+ threadblock_shape=[256, 128, 64],
152
+ stages=3, warp_count=[4, 2, 1],
153
+ math_instruction=math_inst
154
+ )
155
+
156
+ A = TensorDescription(
157
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
158
+ alignment=8
159
+ )
160
+ B = TensorDescription(
161
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
162
+ alignment=8
163
+ )
164
+ C = TensorDescription(
165
+ element=cutlass.float32, layout=cutlass.RowMajor,
166
+ alignment=4
167
+ )
168
+
169
+ element_epilogue = cutlass.float32
170
+
171
+ epilogue_functor = LinearCombination(
172
+ C.element, C.alignment,
173
+ math_inst.element_accumulator, element_epilogue)
174
+
175
+ swizzling_functor = cutlass.IdentitySwizzle1
176
+
177
+ operation = GemmOperationUniversal(
178
+ arch=80, tile_description=tile_description,
179
+ A=A, B=B, C=C,
180
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
181
+ )
182
+
183
+ self.assertTrue(test_all_gemm(operation, "universal"))
184
+
185
+ def test_SM80_Device_Gemm_f16n_f16t_f16t_tensor_op_f16_sliced_k_128x64x64_64x64x32(self):
186
+ math_inst = MathInstruction(
187
+ instruction_shape=[16, 8, 16],
188
+ element_a=cutlass.float16, element_b=cutlass.float16,
189
+ element_accumulator=cutlass.float16, opcode_class=cutlass.OpClass.TensorOp,
190
+ math_operation=MathOperation.multiply_add
191
+ )
192
+
193
+ tile_description = TileDescription(
194
+ threadblock_shape=[128, 64, 64],
195
+ stages=3, warp_count=[2, 1, 1],
196
+ math_instruction=math_inst
197
+ )
198
+
199
+ A = TensorDescription(
200
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
201
+ alignment=8
202
+ )
203
+ B = TensorDescription(
204
+ element=cutlass.float16, layout=cutlass.RowMajor,
205
+ alignment=8
206
+ )
207
+ C = TensorDescription(
208
+ element=cutlass.float16, layout=cutlass.RowMajor,
209
+ alignment=4
210
+ )
211
+
212
+ element_epilogue = cutlass.float16
213
+
214
+ epilogue_functor = LinearCombination(
215
+ C.element, C.alignment,
216
+ math_inst.element_accumulator, element_epilogue)
217
+
218
+ swizzling_functor = cutlass.IdentitySwizzle1
219
+
220
+ operation = GemmOperationUniversal(
221
+ arch=80, tile_description=tile_description,
222
+ A=A, B=B, C=C,
223
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
224
+ )
225
+
226
+ self.assertTrue(test_all_gemm(operation, "universal"))
227
+
228
+ def test_SM80_Device_GemmUniversal_f16n_f16t_f32t_tensor_op_f32_64x64x32_32x32x32(self):
229
+ math_inst = MathInstruction(
230
+ instruction_shape=[16, 8, 16],
231
+ element_a=cutlass.float16, element_b=cutlass.float16,
232
+ element_accumulator=cutlass.float16, opcode_class=cutlass.OpClass.TensorOp,
233
+ math_operation=MathOperation.multiply_add
234
+ )
235
+
236
+ tile_description = TileDescription(
237
+ threadblock_shape=[64, 64, 32],
238
+ stages=10, warp_count=[2, 2, 1],
239
+ math_instruction=math_inst
240
+ )
241
+
242
+ A = TensorDescription(
243
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
244
+ alignment=8
245
+ )
246
+ B = TensorDescription(
247
+ element=cutlass.float16, layout=cutlass.RowMajor,
248
+ alignment=8
249
+ )
250
+ C = TensorDescription(
251
+ element=cutlass.float16, layout=cutlass.RowMajor,
252
+ alignment=4
253
+ )
254
+
255
+ element_epilogue = cutlass.float16
256
+
257
+ epilogue_functor = LinearCombination(
258
+ C.element, C.alignment,
259
+ math_inst.element_accumulator, element_epilogue)
260
+
261
+ swizzling_functor = cutlass.IdentitySwizzle1
262
+
263
+ operation = GemmOperationUniversal(
264
+ arch=80, tile_description=tile_description,
265
+ A=A, B=B, C=C,
266
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
267
+ )
268
+
269
+ self.assertTrue(test_all_gemm(operation, "universal"))
270
+
271
+ def test_SM80_Device_Gemm_f16n_f16t_f32t_tensor_op_f32_256x128x64_64x64x64(self):
272
+ math_inst = MathInstruction(
273
+ instruction_shape=[16, 8, 16],
274
+ element_a=cutlass.float16, element_b=cutlass.float16,
275
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
276
+ math_operation=MathOperation.multiply_add
277
+ )
278
+
279
+ tile_description = TileDescription(
280
+ threadblock_shape=[256, 128, 64],
281
+ stages=3, warp_count=[4, 2, 1],
282
+ math_instruction=math_inst
283
+ )
284
+
285
+ A = TensorDescription(
286
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
287
+ alignment=8
288
+ )
289
+ B = TensorDescription(
290
+ element=cutlass.float16, layout=cutlass.RowMajor,
291
+ alignment=8
292
+ )
293
+ C = TensorDescription(
294
+ element=cutlass.float16, layout=cutlass.RowMajor,
295
+ alignment=8
296
+ )
297
+
298
+ element_epilogue = cutlass.float32
299
+
300
+ epilogue_functor = LinearCombination(
301
+ C.element, C.alignment,
302
+ math_inst.element_accumulator, element_epilogue)
303
+
304
+ swizzling_functor = cutlass.IdentitySwizzle1
305
+
306
+ operation = GemmOperationUniversal(
307
+ arch=80, tile_description=tile_description,
308
+ A=A, B=B, C=C,
309
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
310
+ )
311
+
312
+ self.assertTrue(test_all_gemm(operation, "universal"))
313
+
314
+ def test_test_SM80_Device_Gemm_f16t_f16n_f16t_tensor_op_f16_sliced_k_128x64x64_64x64x32(self):
315
+ math_inst = MathInstruction(
316
+ instruction_shape=[16, 8, 16],
317
+ element_a=cutlass.float16, element_b=cutlass.float16,
318
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
319
+ math_operation=MathOperation.multiply_add
320
+ )
321
+
322
+ tile_description = TileDescription(
323
+ threadblock_shape=[128, 64, 64],
324
+ stages=3, warp_count=[2, 1, 1],
325
+ math_instruction=math_inst
326
+ )
327
+
328
+ A = TensorDescription(
329
+ element=cutlass.float16, layout=cutlass.RowMajor,
330
+ alignment=8
331
+ )
332
+ B = TensorDescription(
333
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
334
+ alignment=8
335
+ )
336
+ C = TensorDescription(
337
+ element=cutlass.float16, layout=cutlass.RowMajor,
338
+ alignment=4
339
+ )
340
+
341
+ element_epilogue = cutlass.float32
342
+
343
+ epilogue_functor = LinearCombination(
344
+ C.element, C.alignment,
345
+ math_inst.element_accumulator, element_epilogue)
346
+
347
+ swizzling_functor = cutlass.IdentitySwizzle1
348
+
349
+ operation = GemmOperationUniversal(
350
+ arch=80, tile_description=tile_description,
351
+ A=A, B=B, C=C,
352
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
353
+ )
354
+
355
+ self.assertTrue(test_all_gemm(operation, "universal"))
356
+
357
+ def test_SM80_Device_Gemm_f16t_f16t_f32n_tensor_op_f32_128x256x64_64x64x64(self):
358
+ math_inst = MathInstruction(
359
+ instruction_shape=[16, 8, 16],
360
+ element_a=cutlass.float16, element_b=cutlass.float16,
361
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
362
+ math_operation=MathOperation.multiply_add
363
+ )
364
+
365
+ tile_description = TileDescription(
366
+ threadblock_shape=[128, 256, 64],
367
+ stages=3, warp_count=[2, 4, 1],
368
+ math_instruction=math_inst
369
+ )
370
+
371
+ A = TensorDescription(
372
+ element=cutlass.float16, layout=cutlass.RowMajor,
373
+ alignment=8
374
+ )
375
+ B = TensorDescription(
376
+ element=cutlass.float16, layout=cutlass.RowMajor,
377
+ alignment=8
378
+ )
379
+ C = TensorDescription(
380
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
381
+ alignment=8
382
+ )
383
+
384
+ element_epilogue = cutlass.float32
385
+
386
+ epilogue_functor = LinearCombination(
387
+ C.element, C.alignment,
388
+ math_inst.element_accumulator, element_epilogue)
389
+
390
+ swizzling_functor = cutlass.IdentitySwizzle1
391
+
392
+ operation = GemmOperationUniversal(
393
+ arch=80, tile_description=tile_description,
394
+ A=A, B=B, C=C,
395
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
396
+ )
397
+
398
+ self.assertTrue(test_all_gemm(operation, "universal"))
399
+
400
+ def test_SM80_Device_Gemm_f16t_f16t_f32t_tensor_op_f32_128x256x64_64x64x64(self):
401
+ math_inst = MathInstruction(
402
+ instruction_shape=[16, 8, 16],
403
+ element_a=cutlass.float16, element_b=cutlass.float16,
404
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
405
+ math_operation=MathOperation.multiply_add
406
+ )
407
+
408
+ tile_description = TileDescription(
409
+ threadblock_shape=[128, 256, 64],
410
+ stages=3, warp_count=[2, 4, 1],
411
+ math_instruction=math_inst
412
+ )
413
+
414
+ A = TensorDescription(
415
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
416
+ alignment=8
417
+ )
418
+ B = TensorDescription(
419
+ element=cutlass.float16, layout=cutlass.ColumnMajor,
420
+ alignment=8
421
+ )
422
+ C = TensorDescription(
423
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
424
+ alignment=4
425
+ )
426
+
427
+ element_epilogue = cutlass.float32
428
+
429
+ epilogue_functor = LinearCombination(
430
+ C.element, C.alignment,
431
+ math_inst.element_accumulator, element_epilogue)
432
+
433
+ swizzling_functor = cutlass.IdentitySwizzle1
434
+
435
+ operation = GemmOperationUniversal(
436
+ arch=80, tile_description=tile_description,
437
+ A=A, B=B, C=C,
438
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
439
+ )
440
+
441
+ self.assertTrue(test_all_gemm(operation, "universal"))
442
+
443
+
444
+
445
+ if __name__ == '__main__':
446
+ pycutlass.get_memory_pool(2**24, 2**24)
447
+ unittest.main()
@@ -0,0 +1,146 @@
1
+ import pycutlass
2
+ from pycutlass import *
3
+ from pycutlass.memory_manager import get_allocated_size
4
+ from pycutlass.test import *
5
+ import unittest
6
+
7
+ from pycutlass.test.gemm_testbed import test_all_gemm
8
+ from pycutlass.utils.device import device_cc
9
+
10
+
11
+ @unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
12
+ class GemmF32nF32nF32nTensorOpF32Sm80(unittest.TestCase):
13
+ def test_SM80_Device_Gemm_f32t_f32n_f32t_tensor_op_bf16_f32_128x128x32_64x64x32(self):
14
+ math_inst = MathInstruction(
15
+ instruction_shape=[16, 8, 8],
16
+ element_a=cutlass.float32, element_b=cutlass.float32,
17
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
18
+ math_operation=MathOperation.multiply_add_fast_bf16
19
+ )
20
+
21
+ tile_description = TileDescription(
22
+ threadblock_shape=[128, 128, 32],
23
+ stages=3, warp_count=[2, 2, 1],
24
+ math_instruction=math_inst
25
+ )
26
+
27
+ A = TensorDescription(
28
+ element=cutlass.float32, layout=cutlass.RowMajor,
29
+ alignment=4
30
+ )
31
+ B = TensorDescription(
32
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
33
+ alignment=4
34
+ )
35
+ C = TensorDescription(
36
+ element=cutlass.float32, layout=cutlass.RowMajor,
37
+ alignment=4
38
+ )
39
+
40
+ element_epilogue = cutlass.float32
41
+
42
+ epilogue_functor = LinearCombination(
43
+ C.element, C.alignment,
44
+ math_inst.element_accumulator, element_epilogue)
45
+
46
+ swizzling_functor = cutlass.IdentitySwizzle1
47
+
48
+ operation = GemmOperationUniversal(
49
+ arch=80, tile_description=tile_description,
50
+ A=A, B=B, C=C,
51
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
52
+ )
53
+
54
+ self.assertTrue(test_all_gemm(operation, "universal"))
55
+
56
+
57
+ def test_SM80_Device_Gemm_f32n_f32n_f32t_tensor_op_f32_128x128x32_64x64x32(self):
58
+ math_inst = MathInstruction(
59
+ instruction_shape=[16, 8, 8],
60
+ element_a=cutlass.float32, element_b=cutlass.float32,
61
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
62
+ math_operation=MathOperation.multiply_add
63
+ )
64
+
65
+ tile_description = TileDescription(
66
+ threadblock_shape=[128, 128, 32],
67
+ stages=3, warp_count=[2, 2, 1],
68
+ math_instruction=math_inst
69
+ )
70
+
71
+ A = TensorDescription(
72
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
73
+ alignment=4
74
+ )
75
+ B = TensorDescription(
76
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
77
+ alignment=4
78
+ )
79
+ C = TensorDescription(
80
+ element=cutlass.float32, layout=cutlass.RowMajor,
81
+ alignment=4
82
+ )
83
+
84
+ element_epilogue = cutlass.float32
85
+
86
+ epilogue_functor = LinearCombination(
87
+ C.element, C.alignment,
88
+ math_inst.element_accumulator, element_epilogue)
89
+
90
+ swizzling_functor = cutlass.IdentitySwizzle1
91
+
92
+ operation = GemmOperationUniversal(
93
+ arch=80, tile_description=tile_description,
94
+ A=A, B=B, C=C,
95
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
96
+ )
97
+
98
+ self.assertTrue(test_all_gemm(operation, "universal"))
99
+
100
+ def test_SM80_Device_Gemm_f32n_f32n_f32t_tensor_op_fast_accurate_f32_64x64x32_32x32x32(self):
101
+ math_inst = MathInstruction(
102
+ instruction_shape=[16, 8, 8],
103
+ element_a=cutlass.float32, element_b=cutlass.float32,
104
+ element_accumulator=cutlass.float32, opcode_class=cutlass.OpClass.TensorOp,
105
+ math_operation=MathOperation.multiply_add_fast_f32
106
+ )
107
+
108
+ tile_description = TileDescription(
109
+ threadblock_shape=[64, 64, 32],
110
+ stages=3, warp_count=[2, 2, 1],
111
+ math_instruction=math_inst
112
+ )
113
+
114
+ A = TensorDescription(
115
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
116
+ alignment=4
117
+ )
118
+ B = TensorDescription(
119
+ element=cutlass.float32, layout=cutlass.ColumnMajor,
120
+ alignment=4
121
+ )
122
+ C = TensorDescription(
123
+ element=cutlass.float32, layout=cutlass.RowMajor,
124
+ alignment=4
125
+ )
126
+
127
+ element_epilogue = cutlass.float32
128
+
129
+ epilogue_functor = LinearCombination(
130
+ C.element, C.alignment,
131
+ math_inst.element_accumulator, element_epilogue)
132
+
133
+ swizzling_functor = cutlass.IdentitySwizzle1
134
+
135
+ operation = GemmOperationUniversal(
136
+ arch=80, tile_description=tile_description,
137
+ A=A, B=B, C=C,
138
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
139
+ )
140
+
141
+ self.assertTrue(test_all_gemm(operation, "universal"))
142
+
143
+ if __name__ == '__main__':
144
+ pycutlass.get_memory_pool(2**24, 2**24)
145
+ pycutlass.compiler.load_from_cache()
146
+ unittest.main()
@@ -0,0 +1,102 @@
1
+ import pycutlass
2
+ from pycutlass import *
3
+ from pycutlass.test import *
4
+ import unittest
5
+
6
+ from pycutlass.test.gemm_testbed import test_all_gemm
7
+ from pycutlass.utils.device import device_cc
8
+
9
+
10
+ @unittest.skipIf(device_cc() < 80, "Device compute capability is insufficient for SM80 tests.")
11
+ class GemmF64TensorOpSm80(unittest.TestCase):
12
+ def test_SM80_Device_Gemm_f64n_f64t_f64t_tensor_op_f64_32x32x16_16x16x16(self):
13
+ math_inst = MathInstruction(
14
+ instruction_shape=[8, 8, 4],
15
+ element_a=cutlass.float64, element_b=cutlass.float64,
16
+ element_accumulator=cutlass.float64, opcode_class=cutlass.OpClass.TensorOp,
17
+ math_operation=MathOperation.multiply_add
18
+ )
19
+
20
+ tile_description = TileDescription(
21
+ threadblock_shape=[32, 32, 16],
22
+ stages=4, warp_count=[2, 2, 1],
23
+ math_instruction=math_inst
24
+ )
25
+
26
+ # alignment 1 restricted for double
27
+ A = TensorDescription(
28
+ element=cutlass.float64, layout=cutlass.ColumnMajor,
29
+ alignment=1
30
+ )
31
+ B = TensorDescription(
32
+ element=cutlass.float64, layout=cutlass.RowMajor,
33
+ alignment=1
34
+ )
35
+ C = TensorDescription(
36
+ element=cutlass.float64, layout=cutlass.RowMajor,
37
+ alignment=1
38
+ )
39
+
40
+ element_epilogue = cutlass.float64
41
+
42
+ epilogue_functor = LinearCombination(
43
+ C.element, C.alignment,
44
+ math_inst.element_accumulator, element_epilogue)
45
+
46
+ swizzling_functor = cutlass.IdentitySwizzle1
47
+
48
+ operation = GemmOperationUniversal(
49
+ arch=80, tile_description=tile_description,
50
+ A=A, B=B, C=C,
51
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
52
+ )
53
+
54
+ self.assertTrue(test_all_gemm(operation, "universal"))
55
+
56
+ def test_SM80_Device_Gemm_f64t_f64n_f64t_tensor_op_f64_64x64x16_32x32x16(self):
57
+ math_inst = MathInstruction(
58
+ instruction_shape=[8, 8, 4],
59
+ element_a=cutlass.float64, element_b=cutlass.float64,
60
+ element_accumulator=cutlass.float64, opcode_class=cutlass.OpClass.TensorOp,
61
+ math_operation=MathOperation.multiply_add
62
+ )
63
+
64
+ tile_description = TileDescription(
65
+ threadblock_shape=[64, 64, 16],
66
+ stages=4, warp_count=[2, 2, 1],
67
+ math_instruction=math_inst
68
+ )
69
+
70
+ # alignment 1 restricted for double
71
+ A = TensorDescription(
72
+ element=cutlass.float64, layout=cutlass.RowMajor,
73
+ alignment=1
74
+ )
75
+ B = TensorDescription(
76
+ element=cutlass.float64, layout=cutlass.ColumnMajor,
77
+ alignment=1
78
+ )
79
+ C = TensorDescription(
80
+ element=cutlass.float64, layout=cutlass.RowMajor,
81
+ alignment=1
82
+ )
83
+
84
+ element_epilogue = cutlass.float64
85
+
86
+ epilogue_functor = LinearCombination(
87
+ C.element, C.alignment,
88
+ math_inst.element_accumulator, element_epilogue)
89
+
90
+ swizzling_functor = cutlass.IdentitySwizzle1
91
+
92
+ operation = GemmOperationUniversal(
93
+ arch=80, tile_description=tile_description,
94
+ A=A, B=B, C=C,
95
+ epilogue_functor=epilogue_functor, swizzling_functor=swizzling_functor
96
+ )
97
+
98
+ self.assertTrue(test_all_gemm(operation, "universal"))
99
+
100
+ if __name__ == '__main__':
101
+ pycutlass.get_memory_pool(2**24, 2**24)
102
+ unittest.main()